mtmd: refactoring: remove clip.h, consolidate code into mtmd (#31865)

Assisted-by: Codex

Signed-off-by: gitGurugu <2049447146@qq.com>

.devops/cpu.Dockerfile

@@ -1,4 +1,4 @@
-ARG UBUNTU_VERSION=24.04
+ARG UBUNTU_VERSION=26.04
 ARG BUILD_DATE=N/A
 ARG APP_VERSION=N/A
 ARG APP_REVISION=N/A
@@ -22,9 +22,7 @@ FROM docker.io/ubuntu:$UBUNTU_VERSION AS build
 ARG TARGETARCH
 
 RUN apt-get update && \
-    apt-get install -y gcc-14 g++-14 build-essential git cmake libssl-dev
-
-ENV CC=gcc-14 CXX=g++-14
+    apt-get install -y build-essential git cmake libssl-dev
 
 WORKDIR /app
 

.devops/vulkan.Dockerfile

@@ -82,20 +82,14 @@ COPY --from=build /app/full /app
 
 WORKDIR /app
 
-ENV PATH="/root/.venv/bin:/root/.local/bin:${PATH}"
-
-# Flag for compatibility with pip
-ARG UV_INDEX_STRATEGY="unsafe-best-match"
 RUN apt-get update \
     && apt-get install -y \
-    build-essential \
-    curl \
     git \
-    ca-certificates \
-    && curl -LsSf https://astral.sh/uv/install.sh | sh \
-    && uv python install 3.13 \
-    && uv venv --python 3.13 /root/.venv \
-    && uv pip install --python /root/.venv/bin/python -r requirements.txt \
+    python3 \
+    python3-pip \
+    python3-wheel \
+    && pip install --break-system-packages --upgrade setuptools \
+    && pip install --break-system-packages -r requirements.txt \
     && apt autoremove -y \
     && apt clean -y \
     && rm -rf /tmp/* /var/tmp/* \

.github/workflows/build-3rd-party.yml

@@ -31,8 +31,8 @@ env:
   LLAMA_ARG_LOG_TIMESTAMPS: 1
 
 jobs:
-  ubuntu-24-llguidance:
-    runs-on: ${{ 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
+  ubuntu-26-llguidance:
+    runs-on: ${{ 'ubuntu-26.04-arm' || 'ubuntu-26.04' }}
 
     steps:
       - name: Clone

.github/workflows/build-and-test-snapdragon.yml

@@ -33,7 +33,7 @@ concurrency:
 
 jobs:
   android-ndk-snapdragon:
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04
     container:
       image: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.8'
     defaults:
@@ -67,7 +67,7 @@ jobs:
           path: pkg-snapdragon/llama.cpp
 
   linux-iot-snapdragon:
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04
     container:
       image: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.8'
     defaults:
@@ -103,7 +103,7 @@ jobs:
   test-snapdragon-qdc:
     name: Test on QDC Device (${{ matrix.device }})
     needs: [android-ndk-snapdragon, linux-iot-snapdragon]
-    runs-on: ubuntu-24.04-arm
+    runs-on: ubuntu-26.04-arm
     timeout-minutes: 180
     strategy:
       fail-fast: false

.github/workflows/build-android.yml

@@ -37,7 +37,7 @@ env:
 
 jobs:
   default:
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04
 
     steps:
       - name: Clone
@@ -56,6 +56,7 @@ jobs:
         uses: android-actions/setup-android@be39fa834029ff78f1a44aa3bb0819b8fc2bd8fd # v4.0.4
         with:
           log-accepted-android-sdk-licenses: false
+          packages: platform-tools cmake;3.31.6
 
       - name: Build
         run: |
@@ -63,7 +64,7 @@ jobs:
           ./gradlew build --no-daemon
 
   ndk:
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04
     container:
       image: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.3'
     defaults:
@@ -101,7 +102,7 @@ jobs:
           path: pkg-adb/llama.cpp
 
   arm64:
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04
 
     env:
       NDK_VERSION: "29.0.14206865"

.github/workflows/build-cache.yml

@@ -40,8 +40,8 @@ jobs:
   #        path: ./spacemit_toolchain
   #        version: ${{ env.SPACEMIT_IME_TOOLCHAIN_VERSION }}
 
-  ubuntu-24-openvino-cache:
-    runs-on: ubuntu-24.04
+  ubuntu-26-openvino-cache:
+    runs-on: ubuntu-26.04
 
     env:
       # Sync versions in build-openvino.yml, ci-self-hosted-openvino.yml, release.yml, release-publish.yml, build-cache.yml, .devops/openvino.Dockerfile, docs/backend/OPENVINO.md

.github/workflows/build-cpu.yml

@@ -50,9 +50,9 @@ jobs:
       matrix:
         include:
           - build: 'x64'
-            os: ubuntu-22.04
+            os: ubuntu-26.04
           - build: 'arm64'
-            os: ubuntu-24.04-arm
+            os: ubuntu-26.04-arm
 
     runs-on: ${{ matrix.os }}
 
@@ -76,13 +76,6 @@ jobs:
             libjpeg-dev build-essential libssl-dev \
             git-lfs
 
-      - name: Toolchain workaround (GCC 14)
-        if: ${{ contains(matrix.os, 'ubuntu-24.04') }}
-        run: |
-          sudo apt-get install -y gcc-14 g++-14
-          echo "CC=gcc-14" >> "$GITHUB_ENV"
-          echo "CXX=g++-14" >> "$GITHUB_ENV"
-
       - name: Python Dependencies
         id: python_depends
         run: |

.github/workflows/build-cross.yml

@@ -163,7 +163,7 @@ jobs:
   #         cmake --build build --config Release -j $(nproc)
 
   debian-13-loongarch64-cpu-cross:
-    runs-on: ${{ 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
+    runs-on: ${{ 'ubuntu-26.04-arm' || 'ubuntu-26.04' }}
     container: debian@sha256:653dfb9f86c3782e8369d5f7d29bb8faba1f4bff9025db46e807fa4c22903671
 
     steps:
@@ -218,7 +218,7 @@ jobs:
           cmake --build build --config Release -j $(nproc)
 
   debian-13-loongarch64-vulkan-cross:
-    runs-on: ${{ 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
+    runs-on: ${{ 'ubuntu-26.04-arm' || 'ubuntu-26.04' }}
     container: debian@sha256:653dfb9f86c3782e8369d5f7d29bb8faba1f4bff9025db46e807fa4c22903671
 
     steps:
@@ -276,8 +276,8 @@ jobs:
 
           cmake --build build --config Release -j $(nproc)
 
-  ubuntu-24-riscv64-cpu-spacemit-ime-cross:
-    runs-on: ubuntu-24.04
+  ubuntu-26-riscv64-cpu-spacemit-ime-cross:
+    runs-on: ubuntu-26.04
 
     env:
       # Make sure this is in sync with build-cache.yml

.github/workflows/build-sycl.yml

@@ -38,7 +38,7 @@ env:
   LLAMA_ARG_LOG_TIMESTAMPS: 1
 
 jobs:
-  ubuntu-24-sycl:
+  ubuntu-26-sycl:
     strategy:
       matrix:
         build: [fp32, fp16]
@@ -48,7 +48,7 @@ jobs:
           - build: fp16
             fp16: ON
 
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-26.04
 
     env:
       ONEAPI_ROOT: /opt/intel/oneapi/
@@ -103,7 +103,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
         with:
-          key: sycl-ubuntu-24-${{ matrix.build }}
+          key: sycl-ubuntu-26-${{ matrix.build }}
           folder: llama.cpp
           hf_bucket: ggml-org/cache
 
@@ -131,7 +131,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
         with:
-          key: sycl-ubuntu-24-${{ matrix.build }}
+          key: sycl-ubuntu-26-${{ matrix.build }}
           folder: llama.cpp
           evict-old-files: 1d
           hf_bucket: ggml-org/cache

.github/workflows/build-virtgpu.yml

@@ -31,8 +31,8 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  ubuntu-24-virtgpu:
-    runs-on: ${{ 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
+  ubuntu-26-virtgpu:
+    runs-on: ${{ 'ubuntu-26.04-arm' || 'ubuntu-26.04' }}
 
     steps:
       - name: Clone

.github/workflows/build-vulkan.yml

@@ -42,7 +42,7 @@ env:
 
 jobs:
   ubuntu-arm64:
-    runs-on: ubuntu-24.04-arm
+    runs-on: ubuntu-26.04-arm
 
     steps:
       - name: Clone
@@ -53,9 +53,7 @@ jobs:
         id: depends
         run: |
           sudo apt-get update
-          sudo apt-get install -y gcc-14 g++-14 build-essential glslc libvulkan-dev spirv-headers libssl-dev ninja-build
-          echo "CC=gcc-14" >> "$GITHUB_ENV"
-          echo "CXX=g++-14" >> "$GITHUB_ENV"
+          sudo apt-get install -y build-essential glslc libvulkan-dev spirv-headers libssl-dev ninja-build
 
       - name: ccache
         uses: ggml-org/ccache-action@v1.2.24
@@ -69,7 +67,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
         with:
-          key: vulkan-ubuntu-24.04-arm
+          key: vulkan-ubuntu-26.04-arm
           folder: llama.cpp
           hf_bucket: ggml-org/cache
 
@@ -92,14 +90,14 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
         with:
-          key: vulkan-ubuntu-24.04-arm
+          key: vulkan-ubuntu-26.04-arm
           folder: llama.cpp
           evict-old-files: 1d
           hf_bucket: ggml-org/cache
           save: true
 
   ubuntu-llvmpipe:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-26.04
 
     steps:
       - name: Clone
@@ -137,7 +135,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
         with:
-          key: vulkan-ubuntu-24.04-llvmpipe
+          key: vulkan-ubuntu-26.04-llvmpipe
           folder: llama.cpp
           hf_bucket: ggml-org/cache
 
@@ -156,7 +154,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
         with:
-          key: vulkan-ubuntu-24.04-llvmpipe
+          key: vulkan-ubuntu-26.04-llvmpipe
           folder: llama.cpp
           evict-old-files: 1d
           hf_bucket: ggml-org/cache

.github/workflows/build-wasi.yml

@@ -49,7 +49,7 @@ env:
 
 jobs:
   ubuntu-webgpu:
-    runs-on: ubuntu-24.04-arm
+    runs-on: ubuntu-26.04-arm
 
     steps:
       - name: Clone
@@ -97,7 +97,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
         with:
-          key: webgpu-ubuntu-24.04-arm-wasi
+          key: webgpu-ubuntu-26.04-arm-wasi
           folder: llama.cpp
           hf_bucket: ggml-org/cache
 
@@ -135,7 +135,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
         with:
-          key: webgpu-ubuntu-24.04-arm-wasi
+          key: webgpu-ubuntu-26.04-arm-wasi
           folder: llama.cpp
           evict-old-files: 1d
           hf_bucket: ggml-org/cache

.github/workflows/build-wasm.yml

@@ -50,7 +50,7 @@ env:
 
 jobs:
   ubuntu-webgpu:
-    runs-on: ubuntu-24.04-arm
+    runs-on: ubuntu-26.04-arm
 
     steps:
       - name: Clone
@@ -84,7 +84,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
         with:
-          key: webgpu-ubuntu-24.04-arm-wasm
+          key: webgpu-ubuntu-26.04-arm-wasm
           folder: llama.cpp
           hf_bucket: ggml-org/cache
 
@@ -107,7 +107,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
         with:
-          key: webgpu-ubuntu-24.04-arm-wasm
+          key: webgpu-ubuntu-26.04-arm-wasm
           folder: llama.cpp
           evict-old-files: 1d
           hf_bucket: ggml-org/cache

.github/workflows/build-webgpu.yml

@@ -34,7 +34,7 @@ env:
 
 jobs:
   format:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-26.04
 
     steps:
       - name: Clone
@@ -45,7 +45,7 @@ jobs:
           wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key |
             sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc > /dev/null
           sudo add-apt-repository -y \
-            "deb http://apt.llvm.org/noble/ llvm-toolchain-noble-22 main"
+            "deb http://apt.llvm.org/resolute/ llvm-toolchain-resolute-22 main"
           sudo apt-get update
           sudo apt-get install -y clang-format-22
 

.github/workflows/copilot-setup-steps.yml

@@ -16,7 +16,7 @@ cache-mode: none
 jobs:
   # The job MUST be called `copilot-setup-steps` or it will not be picked up by Copilot.
   copilot-setup-steps:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-26.04
 
     # Set the permissions to the lowest permissions possible needed for your steps.
     # Copilot will be given its own token for its operations.

.github/workflows/docker.yml

@@ -70,7 +70,7 @@ jobs:
 
   prepare_matrices:
     name: Prepare Docker matrices
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-26.04
     outputs:
       build_matrix: ${{ steps.matrices.outputs.build_matrix }}
       merge_matrix: ${{ steps.matrices.outputs.merge_matrix }}
@@ -87,19 +87,19 @@ jobs:
           # Keep all build targets in one place and derive merge targets from it.
           cat > build-matrix.json <<'JSON'
           [
-            { "tag": "cpu", "dockerfile": ".devops/cpu.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04" },
-            { "tag": "cpu", "dockerfile": ".devops/cpu.Dockerfile", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04-arm" },
-            { "tag": "cpu", "dockerfile": ".devops/s390x.Dockerfile", "platforms": "linux/s390x", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04-s390x", "prebuilt_ui": true },
-            { "tag": "cuda cuda12", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "12.8.1", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" },
-            { "tag": "cuda cuda12", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "12.8.1", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04-arm" },
-            { "tag": "cuda13", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "13.5.0", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" },
-            { "tag": "cuda13", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "13.5.0", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04-arm" },
-            { "tag": "musa", "dockerfile": ".devops/musa.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" },
-            { "tag": "intel", "dockerfile": ".devops/intel.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" },
-            { "tag": "vulkan", "dockerfile": ".devops/vulkan.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04" },
-            { "tag": "vulkan", "dockerfile": ".devops/vulkan.Dockerfile", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04-arm" },
-            { "tag": "rocm", "dockerfile": ".devops/rocm.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" },
-            { "tag": "openvino", "dockerfile": ".devops/openvino.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04" }
+            { "tag": "cpu", "dockerfile": ".devops/cpu.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-26.04" },
+            { "tag": "cpu", "dockerfile": ".devops/cpu.Dockerfile", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-26.04-arm" },
+            { "tag": "cpu", "dockerfile": ".devops/s390x.Dockerfile", "platforms": "linux/s390x", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-26.04-s390x", "prebuilt_ui": true },
+            { "tag": "cuda cuda12", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "12.8.1", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04" },
+            { "tag": "cuda cuda12", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "12.8.1", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04-arm" },
+            { "tag": "cuda13", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "13.5.0", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04" },
+            { "tag": "cuda13", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "13.5.0", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04-arm" },
+            { "tag": "musa", "dockerfile": ".devops/musa.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04" },
+            { "tag": "intel", "dockerfile": ".devops/intel.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04" },
+            { "tag": "vulkan", "dockerfile": ".devops/vulkan.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-26.04" },
+            { "tag": "vulkan", "dockerfile": ".devops/vulkan.Dockerfile", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-26.04-arm" },
+            { "tag": "rocm", "dockerfile": ".devops/rocm.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04" },
+            { "tag": "openvino", "dockerfile": ".devops/openvino.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-26.04" }
           ]
           JSON
 
@@ -400,7 +400,7 @@ jobs:
   merge_arch_tags:
     name: Create shared tags from digests
     needs: [prepare_matrices, push_to_registry, create_tag]
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-26.04
     permissions:
       contents: read
       packages: write

.github/workflows/gguf-publish.yml

@@ -24,7 +24,7 @@ permissions:
 jobs:
   deploy:
 
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04
 
     steps:
     - uses: actions/checkout@v6

.github/workflows/make-release.yml

@@ -39,7 +39,7 @@ permissions:
 
 jobs:
   make-release:
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04
 
     steps:
       - name: Checkout

.github/workflows/python-check-requirements.yml

@@ -24,7 +24,7 @@ concurrency:
 
 jobs:
   python-check-requirements:
-    runs-on: ${{ 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
+    runs-on: ${{ 'ubuntu-26.04-arm' || 'ubuntu-26.04' }}
     name: check-requirements
     steps:
       - name: Check out source repository

.github/workflows/release-publish.yml

@@ -171,7 +171,7 @@ jobs:
             **Linux:**
             - [Ubuntu x64 (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-x64.tar.gz)
             - [Ubuntu arm64 (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-arm64.tar.gz)
-            - [Ubuntu s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz)
+            - [Ubuntu 24.04 s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz)
             - [Ubuntu 26.04 s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-26.04-s390x.tar.gz)
             - [Ubuntu 26.04 s390x (zDNN)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-zdnn-s390x.tar.gz)
             - [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz)

.github/workflows/release.yml

@@ -463,7 +463,7 @@ jobs:
     needs: [check-release, ui-build]
     if: ${{ needs.check-release.outputs.should_release == 'true' }}
 
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-24.04
 
     env:
       NDK_VERSION: "29.0.14206865"
@@ -548,7 +548,7 @@ jobs:
     needs: [check-release, ui-build]
     if: ${{ needs.check-release.outputs.should_release == 'true' }}
 
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     container: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.8'
 
     defaults:
@@ -601,7 +601,7 @@ jobs:
     needs: [check-release, ui-build]
     if: ${{ needs.check-release.outputs.should_release == 'true' }}
 
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     container: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.8'
 
     defaults:

.github/workflows/server.yml

@@ -59,7 +59,7 @@ concurrency:
 
 jobs:
   ubuntu:
-    runs-on: ubuntu-24.04-arm
+    runs-on: ubuntu-26.04-arm
 
     steps:
       - name: Dependencies
@@ -94,7 +94,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
         with:
-          key: server-ubuntu-24.04-arm
+          key: server-ubuntu-26.04-arm
           folder: llama.cpp
           hf_bucket: ggml-org/cache
 
@@ -111,7 +111,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
         with:
-          key: server-ubuntu-24.04-arm
+          key: server-ubuntu-26.04-arm
           folder: llama.cpp
           evict-old-files: 1d
           hf_bucket: ggml-org/cache

.github/workflows/winget.yml

@@ -12,7 +12,7 @@ permissions:
 jobs:
   update:
     name: Update Winget Package
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04
     if: github.repository_owner == 'ggml-org'
 
     steps:

CODEOWNERS

@@ -87,6 +87,7 @@
 /src/llama-graph.*                      @CISC
 /src/llama-model.*                      @CISC
 /src/llama-vocab.*                      @CISC
+/src/llama-moe-cache.*                  @am17an
 /src/models/                            @ggml-org/llama-models
 /tests/                                 @ggerganov
 /tests/test-chat.*                      @pwilkin

common/arg.cpp

@@ -2754,6 +2754,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.moe_cache_size = (size_t) value*1024*1024;
         }
     ).set_env("LLAMA_ARG_MOE_CACHE_MIB"));
+    add_opt(common_arg(
+        {"--moe-cache-cpu-share"}, "F",
+        string_format("share of the MoE cache misses that the CPU computes from the host experts instead of uploading them to the GPU once the cache is full, from 0 to 1 (default: %.2f)", (double) params.moe_cache_cpu_share),
+        [](common_params & params, const std::string & value) {
+            const float share = std::stof(value);
+            if (share < 0.0f || share > 1.0f) {
+                throw std::invalid_argument("invalid value");
+            }
+            params.moe_cache_cpu_share = share;
+        }
+    ).set_env("LLAMA_ARG_MOE_CACHE_CPU_SHARE"));
     add_opt(common_arg(
         {"-ncffn", "--n-cpu-ffn"}, "N",
         "keep the dense FFN weights of the first N layers in the CPU\n"

common/chat-peg-parser.cpp

@@ -348,7 +348,8 @@ void common_chat_peg_mapper::map(const common_peg_ast_node & node) {
     }
 
     if (is_tool_id && current_tool) {
-        auto text = trim_trailing_space(node.text);
+        auto sanitized = node.sanitized_text();
+        auto text      = trim_trailing_space(sanitized);
         if (text.size() >= 2 && text.front() == '"' && text.back() == '"') {
             text = text.substr(1, text.size() - 2);
         }
@@ -401,7 +402,7 @@ void common_chat_peg_mapper::map(const common_peg_ast_node & node) {
     }
 
     if ((is_arg_value || is_arg_string_value) && current_tool) {
-        std::string value_content = std::string(node.text);
+        std::string value_content = node.sanitized_text();
 
         std::string value_to_add;
         if (value_content.empty() && is_arg_string_value) {
@@ -978,7 +979,7 @@ static std::string gemma4_to_json(const common_peg_ast_arena & arena, common_peg
     }
 
     if (node.rule == "gemma4-string-content") {
-        return escape_json_string_inner(std::string(node.text));
+        return escape_json_string_inner(node.sanitized_text());
     }
 
     if (node.rule == "gemma4-string") {
@@ -1195,7 +1196,7 @@ static std::string minimax_m3_value_to_json(const common_peg_ast_arena & arena,
     }
 
     if (node.tag == common_chat_peg_builder::TOOL_ARG_STRING_VALUE) {
-        return "\"" + escape_json_string_inner(std::string(node.text)) + (closed ? "\"" : "");
+        return "\"" + escape_json_string_inner(node.sanitized_text()) + (closed ? "\"" : "");
     }
 
     // Numbers and booleans are written verbatim by the template

common/common.cpp

@@ -1758,7 +1758,8 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.type_k = params.cache_type_k;
     cparams.type_v = params.cache_type_v;
 
-    cparams.moe_cache_size = params.moe_cache_size;
+    cparams.moe_cache_size      = params.moe_cache_size;
+    cparams.moe_cache_cpu_share = params.moe_cache_cpu_share;
 
     return cparams;
 }

common/common.h

@@ -597,7 +597,8 @@ struct common_params {
     ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K
     ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V
 
-    size_t moe_cache_size = 0; // GPU cache size in bytes for the MoE experts kept in the CPU, split among the GPUs like the layers
+    size_t moe_cache_size      = 0;    // GPU cache size in bytes for the MoE experts kept in the CPU, split among the GPUs like the layers
+    float  moe_cache_cpu_share = 0.0f; // share of the MoE cache misses computed on the CPU instead of uploaded
 
     common_conversation_mode conversation_mode = COMMON_CONVERSATION_MODE_AUTO;
 

conversion/__init__.py

@@ -193,6 +193,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "Eagle3LlamaForCausalLM": "llama",
     "LlamaForCausalLMEagle3": "llama",
     "LlavaForConditionalGeneration": "llama",
+    "LlavaOnevision2ForConditionalGeneration": "qwen",
     "LlavaStableLMEpochForCausalLM": "stablelm",
     "MPTForCausalLM": "mpt",
     "MT5ForConditionalGeneration": "t5",
@@ -380,6 +381,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "LightOnOCRForConditionalGeneration": "lighton_ocr",
     "Llama4ForConditionalGeneration": "llama4",
     "LlavaForConditionalGeneration": "llava",
+    "LlavaOnevision2ForConditionalGeneration": "magevl",
     "MERaLiON2ForConditionalGeneration": "ultravox",
     "MiMoV2ForCausalLM": "mimo",
     "MiniMaxM3SparseForConditionalGeneration": "minimax",

conversion/base.py

@@ -3412,6 +3412,9 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st
         arch = text_config["architectures"][0]
     elif model_type == ModelType.MMPROJ and vision_config.get("architectures") is not None:
         arch = vision_config["architectures"][0]
+    # some models (llava_onevision2) only declare the architecture inside text_config
+    if arch is None and text_config.get("architectures") is not None:
+        arch = text_config["architectures"][0]
     if arch is None:
         raise ValueError("Failed to detect model architecture")
     return arch

conversion/chatterbox.py

@@ -355,7 +355,7 @@ class ChatterboxMmprojModel(MmprojModel):
     def set_gguf_parameters(self):
         self.gguf_writer.add_file_type(self.ftype)
 
-        # speaker encoder, the front-end hparams are hard-coded in clip.cpp
+        # speaker encoder, the front-end hparams are hard-coded in tools/mtmd
         self.gguf_writer.add_clip_has_audio_encoder(True)
         self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.CHATTERBOX_SPKENC)
         self.gguf_writer.add_audio_projection_dim(80)

conversion/deepseek.py

@@ -1078,7 +1078,7 @@ class DeepseekV4FlashVisionModel(MmprojModel):
         self.gguf_writer.add_vision_use_silu(True) # SwiGLU MLP
         self.gguf_writer.add_vision_projector_scale_factor(self.hparams_vision["downsample_ratio"])
         self.gguf_writer.add_vision_min_pixels(self.hparams_vision["min_pixels"])
-        # hardcoded on the C++ side (see PROJECTOR_TYPE_DEEPSEEK4V in clip.cpp)
+        # hardcoded on the C++ side (see PROJECTOR_TYPE_DEEPSEEK4V in tools/mtmd)
         # if future models use different values, add GGUF keys for those
         assert self.global_config["vision_max_n_token"] == 384
         assert self.global_config["vision_max_wh_ratio"] == 8

conversion/deepseek41.py

@@ -458,7 +458,7 @@ class DeepseekV41VisionModel(DeepseekV4FlashVisionModel):
         self.gguf_writer.add_vision_use_silu(True)
         self.gguf_writer.add_vision_projector_scale_factor(self.hparams_vision["downsample_ratio"])
         self.gguf_writer.add_vision_min_pixels(self.hparams_vision["min_pixels"])
-        # hardcoded in clip.cpp (PROJECTOR_TYPE_DEEPSEEK41V); add GGUF keys if a later model differs
+        # hardcoded in tools/mtmd (PROJECTOR_TYPE_DEEPSEEK41V); add GGUF keys if a later model differs
         assert self.hparams_vision["max_image_tokens"] == 1024, self.hparams_vision["max_image_tokens"]
         assert self.hparams_vision["max_wh_ratio"] is None, self.hparams_vision["max_wh_ratio"]
         assert self.hparams_vision["rope_theta"] == 10000, self.hparams_vision["rope_theta"]

conversion/drama3.py

@@ -77,7 +77,7 @@ class Drama3MmprojModel(MmprojModel):
         self.gguf_writer.add_gen_audio_head_count(cfg["n_head"])
         self.gguf_writer.add_gen_audio_head_count_kv(cfg["n_local_heads"])
         self.gguf_writer.add_gen_audio_attention_layernorm_eps(cfg["norm_eps"])
-        # other tower and codec hparams are hardcoded in clip.cpp
+        # other tower and codec hparams are hardcoded in tools/mtmd
 
     @classmethod
     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:

conversion/lfm3.py

@@ -181,11 +181,11 @@ class LFM3VLVisionMixin(MmprojModel):
 
     def _set_lfm3vl_gguf_parameters(self):
         assert self.hparams_vision is not None
-        # clip.cpp hardcodes the GELU projector and its input norm eps
+        # mtmd hardcodes the GELU projector and its input norm eps
         if (self.global_config["projector_hidden_act"] != "gelu" or not self.global_config["projector_bias"]
                 or self.global_config["projector_layer_norm_eps"] != 1e-6):
             raise ValueError("unsupported projector_hidden_act / projector_bias / projector_layer_norm_eps")
-        # clip.cpp hardcodes the vision rope base, the RGB patchifier and the 3x3 spatial conv
+        # mtmd hardcodes the vision rope base, the RGB patchifier and the 3x3 spatial conv
         if (self.hparams_vision["rope_parameters"]["rope_theta"] != 1e6 or self.hparams_vision["in_channels"] != 3
                 or self.hparams_vision["conv_kernel_size"] != 3):
             raise ValueError("unsupported vision rope_theta / in_channels / conv_kernel_size")
@@ -347,7 +347,7 @@ class LFM3AudioMmprojModel(MmprojModel):
         "feed_forward.w3":       gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP,
         "feed_forward.w2":       gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN,
     }
-    # the codec graph in clip.cpp hardcodes the liquid-codec FSQ front end and decoder shape, conversion fails on a mismatch
+    # the codec graph in mtmd hardcodes the liquid-codec FSQ front end and decoder shape, conversion fails on a mismatch
     _codec = {
         "model_type": "liquid_codec", "fsq_levels": [8, 7, 6, 6], "n_q": 12, "quantizer_dim": 256,
         "encoder_dim": 512, "stft_n_fft": 1280, "stft_hop_length": 320, "stft_win_length": 1280, "stft_frame_stack": 6,
@@ -405,7 +405,7 @@ class LFM3AudioMmprojModel(MmprojModel):
             if fe[k] != v:
                 raise ValueError(f"unsupported feature_extractor.{k}={fe[k]!r}, expected {v!r}")
 
-        # clip.cpp hardcodes the encoder rope base
+        # mtmd hardcodes the encoder rope base
         if self.hparams_audio["rope_parameters"]["rope_theta"] != 1e6:
             raise ValueError("unsupported audio encoder rope_theta")
 
@@ -420,7 +420,7 @@ class LFM3AudioMmprojModel(MmprojModel):
         if not self._audio_output:
             return
 
-        # audio generation: depthformer hparams, the codec ones are constant and live in clip.cpp
+        # audio generation: depthformer hparams, the codec ones are constant and live in mtmd
         gen = self.get_generation_config()
         for k, v in self._generation.items():
             if gen[k] != v:

conversion/magevl.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from typing import Callable, Iterable, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import MmprojModel, ModelBase, gguf
+
+
+@ModelBase.register("LlavaOnevision2ForConditionalGeneration")
+@ModelBase.example("microsoft/Mage-VL-Cowork")
+class MageVLVisionModel(MmprojModel):
+    # a rope pair uses 2 different angles: the graph runs 2 ropes on heads permuted to [evens | odds]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.hparams_vision is not None
+
+        self.n_embd = self.find_vparam(["hidden_size"])
+        self.n_head = self.find_vparam(["num_attention_heads"])
+        self.d_head = self.n_embd // self.n_head
+
+        d_half = self.d_head // 2
+        assert d_half % 16 == 0, "the 4:6:6 split of the rope requires d_head to be a multiple of 32"
+
+        # 4:6:6 split for t:h:w
+        n_t = 4 * (d_half // 16)
+        n_s = 6 * (d_half // 16)
+        theta = float(self.find_vparam(["rope_theta"]))
+        freq_t = 1.0 / theta ** (torch.arange(n_t, dtype=torch.float64) / n_t)
+        freq_s = 1.0 / theta ** (torch.arange(n_s, dtype=torch.float64) / n_s)
+        freq = torch.cat([freq_t, freq_s, freq_s])
+        self.rope_freq = torch.cat([freq, freq])
+
+        def pair_axis(i: int) -> int:
+            dim = (2 * i) % d_half
+            return 0 if dim < n_t else 1 if dim < n_t + n_s else 2
+
+        self.rope_pairs = sorted(range(d_half), key=pair_axis)
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MAGEVL)
+        self.gguf_writer.add_vision_attention_layernorm_eps(self.find_vparam(["layer_norm_eps"]))
+        self.gguf_writer.add_vision_spatial_merge_size(self.find_vparam(["spatial_merge_size"]))
+        self.gguf_writer.add_vision_min_pixels(int(self.preprocessor_config["min_pixels"]))
+        self.gguf_writer.add_vision_max_pixels(int(self.preprocessor_config["max_pixels"]))
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, _ = item
+
+        if not name.startswith("visual."):
+            return None
+
+        return super().filter_tensors(item)
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # ggml uses theta = pos / rope_freqs
+        pairs = torch.tensor(self.rope_pairs)
+        freq_a = self.rope_freq[2 * pairs]
+        freq_b = self.rope_freq[2 * pairs + 1]
+        yield ("visual.rope_freqs.weight", (1.0 / torch.cat([freq_a, freq_b])).to(torch.float32))
+
+    def permute_qk(self, data_torch: Tensor) -> Tensor:
+        rows = [2 * i for i in self.rope_pairs] + [2 * i + 1 for i in self.rope_pairs]
+        return data_torch.reshape(self.n_head, self.d_head, -1)[:, rows].reshape(data_torch.shape)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if ".self_attn.qkv." in name:
+            q, k, v = data_torch.split(self.n_embd, dim=0)
+            data_torch = torch.cat([self.permute_qk(q), self.permute_qk(k), v], dim=0)
+
+        yield from super().modify_tensors(data_torch, name, bid)

conversion/minicpm.py

@@ -161,7 +161,7 @@ class MiniCPMV4_6VisionModel(MmprojModel):
             # In MiniCPM-V 4.6 `vision_config.image_size` (980) describes the SigLIP
             # positional embedding bucket grid (70 x 70), while the per-slice processing
             # resolution is the preprocessor's `scale_resolution` (typically 448).
-            # The CLIP loader in tools/mtmd/clip.cpp consumes `clip.vision.image_size`
+            # The mmproj loader in tools/mtmd/mtmd-model.cpp consumes `clip.vision.image_size`
             # as the slice size and warmup resolution, so report `scale_resolution` there
             # to match the upstream MiniCPMV4_6ImageProcessorPil slicing rules.
             scale_resolution = self.preprocessor_config.get(
@@ -173,7 +173,7 @@ class MiniCPMV4_6VisionModel(MmprojModel):
         super().set_gguf_parameters()
         assert self.hparams_vision is not None
 
-        # projector type string is consumed by clip_projector_type_from_string() in clip.cpp
+        # projector type string is consumed by projector_type_from_string() in tools/mtmd/mtmd-impl.h
         self.gguf_writer.add_clip_projector_type(self.projector_type)
 
         self.gguf_writer.add_vision_projector_scale_factor(
@@ -243,7 +243,7 @@ class MiniCPMV4_7VisionModel(MiniCPMV4_6VisionModel):
     projector_type = gguf.VisionProjectorType.MINICPMV4_7
     # MiniCPMV4_7ImageProcessorPil default
     default_scale_resolution = 448
-    # rows of v.tok_embd_sep, the order must match clip_suffix_rows() in clip-impl.h
+    # rows of v.tok_embd_sep, the order must match mtmd_suffix_rows() in tools/mtmd/mtmd-impl.h
     tok_embd_sep = ["</image>", "<slice>", "</slice>", "\n"]
 
     def get_downsample_mode(self) -> str:

conversion/nemotron.py

@@ -208,7 +208,7 @@ class NemotronNanoV2VLModel(MmprojModel):
         return name, gen
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # RADIO's pos_embed doesn't have .weight suffix, but clip.cpp expects it
+        # RADIO's pos_embed doesn't have .weight suffix, but mtmd expects it
         if "patch_generator.pos_embed" in name:
             # Downsample position embeddings for fixed 512x512 image size
             import torch.nn.functional as F

conversion/pockettts.py

@@ -302,7 +302,7 @@ class PocketTTSMmprojModel(MmprojModel):
         # voice-prompt encoder: mimi encoder + speaker_proj
         if self.has_audio_encoder:
             self.gguf_writer.add_clip_has_audio_encoder(True)
-            # note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models
+            # note: the 24kHz sample rate is hardcoded on the mtmd side, like the other audio models
             self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC)
             self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
             self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"])
@@ -314,7 +314,7 @@ class PocketTTSMmprojModel(MmprojModel):
             self.gguf_writer.add_audio_num_mel_bins(1)
 
         # generation: flow-matching decoder + mimi decoder
-        # the SEANet and flow net hparams are constant across the family, clip.cpp holds them
+        # the SEANet and flow net hparams are constant across the family, mtmd holds them
         self.gguf_writer.add_clip_has_gen_audio_encoder(True)
         self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_GEN)
         self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text)

conversion/qwen.py

@@ -170,7 +170,7 @@ class Qwen2MoeModel(TextModel):
                 raise ValueError(f"Unprocessed experts: {experts}")
 
 
-@ModelBase.register("Qwen3ForCausalLM", "Qwen3Model")
+@ModelBase.register("Qwen3ForCausalLM", "Qwen3Model", "LlavaOnevision2ForConditionalGeneration")
 @ModelBase.example("Qwen/Qwen3-8B")
 class Qwen3Model(Qwen2Model):
     model_arch = gguf.MODEL_ARCH.QWEN3

conversion/qwen3tts.py

@@ -235,11 +235,11 @@ class Qwen3TTSSpeakerEncoderModel(MmprojModel):
 
         # handle speaker encoder config
         self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
-        # mel_spectrogram() front-end: sr=24000, n_fft=1024, hop=256, n_mels=128, fmin=0, fmax=12000 (=sr/2, the clip.cpp default)
+        # mel_spectrogram() front-end: sr=24000, n_fft=1024, hop=256, n_mels=128, fmin=0, fmax=12000 (=sr/2, the mtmd default)
         self.gguf_writer.add_audio_num_mel_bins(128)
         # 3 SE-Res2Net stages; the stem conv, mfa, asp and fc are not counted here
         self.gguf_writer.add_audio_block_count(3)
-        # ECAPA-TDNN has no attention/FFN, these are dummy to allow clip.cpp to load it
+        # ECAPA-TDNN has no attention/FFN, these are dummy to allow mtmd to load it
         self.gguf_writer.add_audio_embedding_length(1536)
         self.gguf_writer.add_audio_head_count(1)
         self.gguf_writer.add_audio_feed_forward_length(1536)
@@ -256,7 +256,7 @@ class Qwen3TTSSpeakerEncoderModel(MmprojModel):
         self.gguf_writer.add_gen_audio_head_count(code_predictor_config["num_attention_heads"])
         self.gguf_writer.add_gen_audio_head_count_kv(code_predictor_config["num_key_value_heads"])
         self.gguf_writer.add_gen_audio_attention_layernorm_eps(code_predictor_config["rms_norm_eps"])
-        # note: code2wav hparams are hardcoded on the mtmd/clip.cpp side for now, not written here
+        # note: code2wav hparams are hardcoded on the mtmd side for now, not written here
 
     def _wav_decoder_config(self) -> dict[str, Any] | None:
         # code2wav has its own config.json, inside the speech_tokenizer dir

docs/backend/SYCL.md

@@ -817,7 +817,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_ENABLE_VMM | 0 or 1 (default) | Enable the virtual-memory device pool. |
 | GGML_SYCL_ENABLE_MKL_FA | 1 (default) or 0 | Enable oneMKL GEMM flash attention for XMX-accelerated prompt processing with quantized KV cache. Automatically activates during prefill (prompt processing) when all conditions are met: (1) flash-attn enabled (`-fa` or `--flash-attn on`), (2) KV cache quantized (`--cache-type-k q8_0 --cache-type-v q8_0` or other `*_0/*_1` types), (3) batch size ≥ 1024 (`--batch-size 1024`), (4) prompt length ≥ 1024 tokens. Set to 0 to force the TILE kernel for A/B testing. Example minimum command: `llama-cli -m model.gguf -fa -ngl 99 --cache-type-k q8_0 --cache-type-v q8_0 --batch-size 1024 -p "your prompt"` |
 | GGML_SYCL_MKL_FA_DEBUG | 0 (default) or 1 | Enable per-call diagnostic logging for MKL flash attention: GEMM/softmax timings, interleaved-head detection, and buffer memory usage. |
-| GGML_SYCL_MEMTRACE | 0 (default), 1, 2 | Enable record and output memory allocation diagnostics. Requires `-lv 4`. <br>0 -  Disable<br>1 - Basic memory info, including current and peak allocations, as well allocations from other sources, around 50 lines per model load.<br>2 - More verbose, logging around 900 specific allocations and deallocations. |
+| GGML_SYCL_MEMTRACE | 0 (default), 1, 2 | Enable record and output memory allocation diagnostics. Requires `-lv 4`. <br>0 -  Disable<br>1 - Basic memory info, including current and peak allocations, as well allocations from other sources, and the device free memory before each model buffer allocation, around 50 lines per model load.<br>2 - More verbose, logging around 900 specific allocations and deallocations. |
 | GGML_SYCL_MEMTRACE_STEP | 64 (default) or positive integer | With GGML_SYCL_MEMTRACE=1, the minimum growth in memory usage to trigger another log record. |
 | GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. |
 | GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. Unsupported types and layouts fall back to the standalone op kernels. See `ggml_sycl_can_fuse()`. |
@@ -836,6 +836,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | ZES_ENABLE_SYSMAN | 0 (default) or 1 | Support to get free memory of GPU by sycl::aspect::ext_intel_free_memory.<br>Recommended to use when --split-mode = layer |
 | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. |
 | GGML_SYCL_USM_SYSTEM | 0 (default) or 1 | Enable experimental support for [USM system allocations](https://github.khronos.org/SYCL_Reference/iface/usm_basic_concept.html#system-allocations) for large GPU buffers. This requires enough host memory for model weights and caches, an Intel Xe2+ GPU such as BMG or newer and supported on Linux only, with CONFIG_DRM_XE_GPUSVM enabled. |
+| GGML_SYCL_BMG_VRAM_CHECK | 1 (default) or 0 | Workaround for BMG (Arc B-series) GPUs, where the Level Zero driver accepts device allocations beyond VRAM and moves the overflow to system RAM, which can exhaust host memory when a model does not fit (e.g. `-ngl 999`, see https://github.com/intel/compute-runtime/issues/1047). When enabled, a model or split buffer that does not fit in the device's free memory is refused and the model load fails with an error. Set to 0 to disable. |
 
 1. <a id="sycl-graph-footnote"></a>SYCL Graph support (`GGML_SYCL_ENABLE_GRAPH=1`) is limited to single-GPU workloads and may also depend on the model architecture and kernels in use.
    Performance depends on how often a recorded graph can be reused, which will vary by workload.

docs/backend/snapdragon/developer.md

@@ -255,7 +255,7 @@ Writing high-performance operators for Hexagon requires following specific guide
 ### Tracing Instrumentation
 
 - All kernels must include trace events for performance profiling and timeline visualization in Perfetto
-  ([`hex-profile.h`](../../../ggml/src/ggml-hexagon/htp/hex-profile.h)).
+  ([`htp-trace.h`](../../../ggml/src/ggml-hexagon/htp/htp-trace.h)).
 - Surround compute sections with `htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) info)` and
   `htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) info)`.
 - Use specific event types for major phases:

docs/development/HOWTO-add-model.md

@@ -128,7 +128,7 @@ Note: to debug the inference graph: you can use [llama-eval-callback](/examples/
 If the new model supports multimodal inputs, you will need to add a new encoder definition in `libmtmd`. You can find more information about llama.cpp's multimodal support in [the docs](../multimodal.md) and in the `tools/mtmd` source directory.
 
 1. In the conversion script, make sure you add a subclass that extends `MmprojModel` or another class that inherits from the same base class.
-2. Add the encoder definition in `clip.cpp`.
+2. Add the encoder definition: hparams and tensor loading in `mtmd-model.cpp`, graph inputs in `mtmd-subcontext.cpp`.
 3. Implement the preprocessor in `mtmd.cpp`. In most cases, you can reuse an existing preprocessor.
 4. Implement the encoder GGML graph, either in a dedicated file if the model is truly different from existing ones, or by reusing an existing implementation (for example: siglip, pixtral, or qwen) and adding a model-specific projector.
 

docs/moe-cache.md

@@ -0,0 +1,26 @@
+# MoE cache
+
+When the experts of MoE layers stay in host memory (`-cmoe`, `-ncmoe`, `-ot`), the MoE cache keeps the most used ones in device memory for small batches (decode, speculative decoding). Large batches still upload the experts they use.
+
+| Argument | |
+|---|---|
+| `--moe-cache-mib N` | cache size in MiB, split among the GPUs like the layers (default: 0, disabled) |
+| `--moe-cache-cpu-share F` | share of the cache misses computed on the CPU from the host experts instead of uploaded, once the cache is full (default: 0) |
+
+The cache hits run on the device. With `--moe-cache-cpu-share F` the CPU computes a share `F` of the misses while the device computes the hits and uploads the other misses into the cache. This is faster when the CPU reads the host experts faster than the device can upload them, for example with fast RAM, a PCIe 4 link or experts that are not in pinned memory. The host experts are pinned with `-lm none`, `-lm mlock` and `-lm dio`, unless `--no-host`, but not when they are mmap'd (the default). The device uploads pinned experts at the full speed of the link, so the CPU share often does not help then. Until the cache is full every miss is uploaded, so the cache fills as fast as without the CPU share. The best share depends on the machine, the model and the cache size.
+
+## Picking the arguments
+
+Put the script in front of your server command:
+
+```bash
+python3 scripts/moe-cache-tune.py -- build/bin/llama-server -m model.gguf -cmoe -t 8 -c 32768
+```
+
+It starts the server 6 times, at two cache sizes × shares 0, 0.5 and 0.95, and times the generation of a few thousand tokens. That takes a few minutes, mostly model loading.
+
+From these runs it fits the upload and CPU speeds of the machine, then replays the routing of all the runs through a copy of the cache logic for other sizes and shares. It prints the measured speeds, the predicted t/s and the recommended arguments: the lowest share within 5% of the fastest.
+
+- **Server command:** keep your own arguments (`-t`, `-c`, expert placement, load mode); the speeds and the free device memory depend on them.
+- **`--max-mib N`:** tunes a smaller cache than what fits on the device. The runs in `--out` are reused while the command stays the same, so trying another value is instant.
+- **`--markdown`:** prints Markdown tables to paste in a PR.

docs/ops.md

@@ -34,7 +34,7 @@ Legend:
 |                CONV_TRANSPOSE_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                CONV_TRANSPOSE_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              COS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                      COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                      COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              CPY | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
 |               CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |

docs/ops/WebGPU.csv

@@ -5465,8 +5465,10 @@
 "WebGPU: WebGPU","CONV_TRANSPOSE_2D","kernel_type=f16,ne_input=[3,2,3,1],ne_kernel=[2,2,1,3],stride=1","support","0","no","WebGPU"
 "WebGPU: WebGPU","CONV_TRANSPOSE_2D","kernel_type=f16,ne_input=[10,10,9,1],ne_kernel=[3,3,1,9],stride=2","support","0","no","WebGPU"
 "WebGPU: WebGPU","CONV_TRANSPOSE_2D","kernel_type=f16,ne_input=[129,63,35,1],ne_kernel=[3,3,48,35],stride=1","support","0","no","WebGPU"
-"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,500,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,5000,1,1]","support","0","no","WebGPU"
+"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,500,1,1],same_src=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,5000,1,1],same_src=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,500,1,1],same_src=1","support","1","yes","WebGPU"
+"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,5000,1,1],same_src=1","support","1","yes","WebGPU"
 "WebGPU: WebGPU","ARGMAX","type=f32,ne=[32,1,1,1]","support","1","yes","WebGPU"
 "WebGPU: WebGPU","ARGMAX","type=f32,ne=[32,513,1,1]","support","1","yes","WebGPU"
 "WebGPU: WebGPU","ARGMAX","type=f32,ne=[100,10,1,1]","support","1","yes","WebGPU"

examples/speculative/speculative.cpp

@@ -103,6 +103,9 @@ int main(int argc, char ** argv) {
     params.cpuparams_batch.n_threads = params.speculative.draft.cpuparams_batch.n_threads;
     params.tensor_buft_overrides     = params.speculative.draft.tensor_buft_overrides;
 
+    // the MoE cache is only used by the target context
+    params.moe_cache_size = 0;
+
     auto llama_init_dft = common_init_from_params(params);
 
     model_dft = llama_init_dft->model();

ggml/include/ggml.h

@@ -1519,6 +1519,7 @@ extern "C" {
     // note: this op has undefined behaviour if a given `a` is referenced more than once for a given `b`
     //       (i.e. duplicate experts per token are not guaranteed to be supported in all backends)
     // TODO: add backend support: https://github.com/ggml-org/llama.cpp/pull/29998#issuecomment-6153742533
+    // the CPU backend skips negative ids, their rows of the result are zero
     GGML_API struct ggml_tensor * ggml_mul_mat_id(
             struct ggml_context * ctx,
             struct ggml_tensor  * as,

ggml/src/CMakeLists.txt

@@ -216,6 +216,12 @@ set_target_properties(ggml-base PROPERTIES
 configure_file(ggml-version.h.in ${CMAKE_CURRENT_BINARY_DIR}/ggml-version.h @ONLY)
 
 target_include_directories(ggml-base PRIVATE . ${CMAKE_CURRENT_BINARY_DIR})
+
+# icx/icpx default to -fp-model=fast, which contracts and reassociates the reference quantizers, so their output
+# depends on the compiler; keep ggml-base IEEE-exact so the CPU reference matches gcc/clang builds
+if (CMAKE_C_COMPILER_ID STREQUAL "IntelLLVM")
+    target_compile_options(ggml-base PRIVATE -fp-model=precise)
+endif()
 if (GGML_BACKEND_DL)
     target_compile_definitions(ggml-base PUBLIC GGML_BACKEND_DL)
 endif()

ggml/src/ggml-backend.cpp

@@ -1246,6 +1246,14 @@ struct ggml_backend_sched {
     int prev_graph_size;
 
     int debug_realloc;
+
+    // inputs of the next split read from the backend of the current split before it computes, see ggml_backend_sched_prefetch_inputs
+    ggml_backend_event_t prefetch_events[GGML_SCHED_MAX_BACKENDS]; // created on first use, NULL if the backend has no events
+    bool                 prefetch_tried [GGML_SCHED_MAX_BACKENDS];
+    struct ggml_tensor * prefetched[GGML_SCHED_MAX_SPLIT_INPUTS];
+    int n_prefetched;
+    int prefetch_split;     // the split that reads the prefetched inputs
+    int prefetch_backend_id; // the backend that copies them
 };
 
 #define hash_id(tensor) ggml_hash_find_or_insert(&sched->hash_set, tensor)
@@ -2086,6 +2094,16 @@ static void ggml_backend_sched_copy_input(ggml_backend_sched_t sched, struct ggm
     ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, input);
     struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy);
 
+    // read before the previous split, wait for the copy only
+    if (split == &sched->splits[sched->prefetch_split]) {
+        for (int i = 0; i < sched->n_prefetched; i++) {
+            if (sched->prefetched[i] == input) {
+                ggml_backend_event_synchronize(sched->prefetch_events[sched->prefetch_backend_id]);
+                return;
+            }
+        }
+    }
+
     if (input->flags & GGML_TENSOR_FLAG_INPUT) {
         // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done
         if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
@@ -2122,6 +2140,62 @@ static void ggml_backend_sched_copy_input(ggml_backend_sched_t sched, struct ggm
     }
 }
 
+// the inputs of the next split that are computed on the backend of this split before it, and copied to host memory,
+// are copied before this split computes, so that the next split waits for these copies instead of this split
+// e.g. a CPU split runs while a GPU split computes, if it reads only the results of the earlier splits
+static void ggml_backend_sched_prefetch_inputs(ggml_backend_sched_t sched, int split_id) {
+    sched->n_prefetched = 0;
+    if (split_id + 1 >= sched->n_splits) {
+        return;
+    }
+    struct ggml_backend_sched_split * split = &sched->splits[split_id];
+    struct ggml_backend_sched_split * next  = &sched->splits[split_id + 1];
+    const int backend_id = split->backend_id;
+    ggml_backend_t backend = sched->backends[backend_id];
+    if (backend->iface.get_tensor_async == NULL) {
+        return;
+    }
+
+    for (int i = 0; i < next->n_inputs && sched->n_prefetched < GGML_SCHED_MAX_SPLIT_INPUTS; i++) {
+        struct ggml_tensor * input = next->inputs[i];
+        if ((input->flags & GGML_TENSOR_FLAG_INPUT) || tensor_backend_id(input) != backend_id) {
+            continue;
+        }
+        struct ggml_tensor * input_cpy = tensor_copy(input, next->backend_id, sched->cur_copy);
+        if (input_cpy->buffer == NULL || !ggml_backend_buffer_is_host(input_cpy->buffer)) {
+            continue;
+        }
+
+        // not if this split computes or writes the input
+        const struct ggml_tensor * base = input->view_src ? input->view_src : input;
+        bool written = false;
+        for (int j = 0; j < split->graph.n_nodes && !written; j++) {
+            const struct ggml_tensor * node = split->graph.nodes[j];
+            written = !ggml_is_view_op(node->op) && (node == base || node->view_src == base);
+        }
+        if (written) {
+            continue;
+        }
+
+        if (!sched->prefetch_tried[backend_id]) {
+            sched->prefetch_tried[backend_id]  = true;
+            sched->prefetch_events[backend_id] = ggml_backend_event_new(backend->device);
+        }
+        if (sched->prefetch_events[backend_id] == NULL) {
+            return;
+        }
+
+        ggml_backend_tensor_get_async(backend, input, input_cpy->data, 0, ggml_nbytes(input));
+        sched->prefetched[sched->n_prefetched++] = input;
+    }
+
+    if (sched->n_prefetched > 0) {
+        sched->prefetch_split      = split_id + 1;
+        sched->prefetch_backend_id = backend_id;
+        ggml_backend_event_record(sched->prefetch_events[backend_id], backend);
+    }
+}
+
 static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) {
     GGML_ASSERT(sched);
     struct ggml_backend_sched_split * splits = sched->splits;
@@ -2150,6 +2224,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                 ggml_backend_sched_copy_input(sched, split, split->inputs[input_id]);
             }
         }
+        // after the copies above, which may read host memory reused by the copies of the next split
+        // before the weights, whose copy callback may queue large uploads
+        ggml_backend_sched_prefetch_inputs(sched, split_id);
         for (int input_id = 0; input_id < split->n_inputs; input_id++) {
             if (ggml_backend_sched_is_host_weight(split->inputs[input_id])) {
                 ggml_backend_sched_copy_input(sched, split, split->inputs[input_id]);
@@ -2286,6 +2363,7 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
         for (int c = 0; c < sched->n_copies; c++) {
             ggml_backend_event_free(sched->events[b][c]);
         }
+        ggml_backend_event_free(sched->prefetch_events[b]);
     }
     ggml_gallocr_free(sched->galloc);
     ggml_free(sched->ctx);

ggml/src/ggml-cpu/ggml-cpu.c

@@ -224,8 +224,8 @@ static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = {
     },
     [GGML_TYPE_F16] = {
         .from_float               = (ggml_from_float_t) ggml_cpu_fp32_to_fp16,
-        .vec_dot                  = (ggml_vec_dot_t) ggml_vec_dot_f16,
-        .vec_dot_type             = GGML_TYPE_F16,
+        .vec_dot                  = (ggml_vec_dot_t) ggml_vec_dot_f16_f32,
+        .vec_dot_type             = GGML_TYPE_F32,
         .nrows                    = 1,
     },
     [GGML_TYPE_Q1_0] = {
@@ -1284,9 +1284,12 @@ void ggml_compute_forward_mul_mat(
     const int ith = params->ith;
     const int nth = params->nth;
 
-    enum ggml_type           const vec_dot_type         = type_traits_cpu[src0->type].vec_dot_type;
-    ggml_from_float_t        const from_float           = type_traits_cpu[vec_dot_type].from_float;
-    int64_t                  const vec_dot_num_rows     = type_traits_cpu[src0->type].nrows;
+    enum ggml_type    const vec_dot_type     = type_traits_cpu[src0->type].vec_dot_type;
+    ggml_from_float_t const from_float       = type_traits_cpu[vec_dot_type].from_float;
+    int64_t           const vec_dot_num_rows = type_traits_cpu[src0->type].nrows;
+
+    // TODO: a more efficient F16 x F16 path can be used here if ggml_prec_get_src(dst, 1) allows converting src1 to F16
+    //       the code needs to be refactored to enable this (src1 is currently converted based on vec_dot_type alone)
 
     GGML_ASSERT(ne0 == ne01);
     GGML_ASSERT(ne1 == ne11);
@@ -1509,8 +1512,8 @@ static void ggml_compute_forward_mul_mat_id_one_chunk(
 
     const enum ggml_type type = src0->type;
 
-    ggml_vec_dot_t    const vec_dot      = type_traits_cpu[type].vec_dot;
-    enum ggml_type    const vec_dot_type = type_traits_cpu[type].vec_dot_type;
+    ggml_vec_dot_t const vec_dot      = type_traits_cpu[type].vec_dot;
+    enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type;
 
     const int64_t blck_0 = 16;
     const int64_t blck_1 = 16;
@@ -1577,8 +1580,11 @@ static void ggml_compute_forward_mul_mat_id(
 
     const bool src1_cont = ggml_is_contiguous(src1);
 
-    enum ggml_type    const vec_dot_type    = type_traits_cpu[type].vec_dot_type;
-    ggml_from_float_t const from_float      = type_traits_cpu[vec_dot_type].from_float;
+    enum ggml_type    const vec_dot_type = type_traits_cpu[type].vec_dot_type;
+    ggml_from_float_t const from_float   = type_traits_cpu[vec_dot_type].from_float;
+
+    // TODO: a more efficient F16 x F16 path can be used here if ggml_prec_get_src(dst, 1) allows converting src1 to F16
+    //       the code needs to be refactored to enable this (src1 is currently converted based on vec_dot_type alone)
 
     // we don't support permuted src0 or src1
     GGML_ASSERT(nb00 == ggml_type_size(type));
@@ -2965,8 +2971,8 @@ struct ggml_cplan ggml_graph_plan(
                         if ((node->src[0]->type == GGML_TYPE_F16 ||
                              node->src[0]->type == GGML_TYPE_BF16) &&
                             node->src[1]->type == GGML_TYPE_F32) {
+                            cur += sizeof(float)*ne10*ne11;
                             cur += sizeof(ggml_fp16_t)*ne00*ne01*ne02;
-                            cur += sizeof(ggml_fp16_t)*ne10*ne11;
                         } else if (node->src[0]->type == GGML_TYPE_F32 &&
                                    node->src[1]->type == GGML_TYPE_F32) {
                             cur += sizeof(float)*ne00*ne01*ne02;
@@ -2995,8 +3001,8 @@ struct ggml_cplan ggml_graph_plan(
                         GGML_ASSERT(node->src[0]->type == GGML_TYPE_F16 || node->src[0]->type == GGML_TYPE_F32);
                         GGML_ASSERT(node->src[1]->type == GGML_TYPE_F32);
 
+                        cur += sizeof(float) * ne10 * ne11 * ne12 * ne13;
                         cur += ggml_type_size(node->src[0]->type) * ne00 * ne01 * ne02 * ne03;
-                        cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12 * ne13;
 
                     } break;
                 case GGML_OP_TOP_K:

ggml/src/ggml-cpu/ggml-cpu.cpp

@@ -455,10 +455,11 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
                 src0->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32) {
                 return src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16;
             }
-            // BF16 in src1 is widened into the F32 work buffer
+            // F16 or BF16 in src1 is widened into the F32 work buffer
             return src1->type == GGML_TYPE_F32 ||
                    src1->type == ggml_get_type_traits_cpu(src0->type)->vec_dot_type ||
-                   (src1->type == GGML_TYPE_BF16 && ggml_get_type_traits_cpu(src0->type)->vec_dot_type == GGML_TYPE_F32);
+                   ((src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_BF16) &&
+                    ggml_get_type_traits_cpu(src0->type)->vec_dot_type == GGML_TYPE_F32);
         case GGML_OP_SOFT_MAX_BACK: {
             if (op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32) {
                 return false;

ggml/src/ggml-cpu/llamafile/sgemm.cpp

@@ -326,6 +326,12 @@ template <> inline float32x4_t load(const ggml_fp16_t * p) {
 
     return vec_xl(0, (const float *)(tmp));
 }
+template <> inline float32x4_t load(const ggml_bf16_t * p) {
+    uint64_t bits;
+    memcpy(&bits, p, sizeof(bits));
+    const uint64x2_t v_bits = { bits, bits };
+    return GGML_BF16_TO_F32_HI((GGML_BF16_VEC) v_bits);
+}
 template <> inline float32x4_t load(const float * p) {
     return vec_xl(0, p);
 }
@@ -1302,6 +1308,88 @@ class tinyBLAS_RVV {
 //////////////////////////////////////////////////////////////////////////////////////////
 // QUANT ZERO MATRIX MULTIPLICATION
 
+#if defined(__VXE__) || defined(__VXE2__)
+template <typename TA>
+class tinyBLAS_Q0_VXE {
+    static_assert(sizeof(TA::qs) == QK8_0, "VXE weight loads require 32 quantized bytes per block");
+
+  public:
+    tinyBLAS_Q0_VXE(int64_t k,
+                    const TA * A, int64_t lda,
+                    const block_q8_0 * B, int64_t ldb,
+                    float * C, int64_t ldc,
+                    int32_t ith, int32_t nth)
+        : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) {
+    }
+
+    void matmul(int64_t m, int64_t n) {
+        constexpr int32_t tile_size = 2;
+        const int64_t row_tiles = m / tile_size;
+        const int64_t tiles     = row_tiles * (n / tile_size);
+        const int64_t duty      = (tiles + nth - 1) / nth;
+        const int64_t start     = duty * ith;
+        const int64_t end       = MIN(start + duty, tiles);
+
+        for (int64_t job = start; job < end; ++job) {
+            const int64_t ii = tile_size * (job % row_tiles);
+            const int64_t jj = tile_size * (job / row_tiles);
+            float32x4_t v_sum[tile_size][tile_size] = {};
+
+            for (int64_t l = 0; l < k; ++l) {
+                int8x16_t v_xl[tile_size], v_xh[tile_size], v_yl[tile_size], v_yh[tile_size];
+                float x_scale[tile_size], y_scale[tile_size];
+
+                for (int32_t i = 0; i < tile_size; ++i) {
+                    const TA & x = A[lda * (ii + i) + l];
+                    const block_q8_0 & y = B[ldb * (jj + i) + l];
+
+                    const int8_t * x_qs = x.qs;
+                    const int8_t * y_qs = y.qs;
+                    v_xl[i] = vec_xl(0, x_qs);
+                    v_xh[i] = vec_xl(0, x_qs + QK8_0 / 2);
+                    v_yl[i] = vec_xl(0, y_qs);
+                    v_yh[i] = vec_xl(0, y_qs + QK8_0 / 2);
+                    x_scale[i] = unhalf(x.d);
+                    y_scale[i] = unhalf(y.d);
+                }
+
+                for (int32_t j = 0; j < tile_size; ++j) {
+                    for (int32_t i = 0; i < tile_size; ++i) {
+                        const int32x4_t v_dot = dot_s8_wide(v_xl[i], v_yl[j]) + dot_s8_wide(v_xh[i], v_yh[j]);
+                        const float32x4_t v_scale = vec_splats(x_scale[i] * y_scale[j]);
+                        v_sum[j][i] = vec_madd(vec_float(v_dot), v_scale, v_sum[j][i]);
+                    }
+                }
+            }
+
+            for (int32_t j = 0; j < tile_size; ++j) {
+                for (int32_t i = 0; i < tile_size; ++i) {
+                    C[ldc * (jj + j) + ii + i] = hsum(v_sum[j][i]);
+                }
+            }
+        }
+    }
+
+  private:
+    static int32x4_t dot_s8_wide(int8x16_t v_x, int8x16_t v_y) {
+        const int16x8_t v_even = vec_mule(v_x, v_y);
+        const int16x8_t v_odd  = vec_mulo(v_x, v_y);
+        // Two (-128) * (-128) products can overflow a signed 16-bit sum.
+        return vec_unpackh(v_even) + vec_unpackl(v_even) + vec_unpackh(v_odd) + vec_unpackl(v_odd);
+    }
+
+    const TA * const A;
+    const block_q8_0 * const B;
+    float * const C;
+    const int64_t k;
+    const int64_t lda;
+    const int64_t ldb;
+    const int64_t ldc;
+    const int32_t ith;
+    const int32_t nth;
+};
+#endif
+
 #if defined(__ARM_FEATURE_DOTPROD)
 template <typename TA>
 class tinyBLAS_Q0_ARM {
@@ -4006,6 +4094,14 @@ bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64
                 (float *)C, ldc};
             return tb.matmul(m, n);
         }
+#elif defined(__VXE__) || defined(__VXE2__)
+        if (Btype == GGML_TYPE_BF16) {
+            tinyBLAS<4, float32x4_t, float32x4_t, ggml_bf16_t, ggml_bf16_t, float> tb{ params, k,
+                (const ggml_bf16_t *)A, lda,
+                (const ggml_bf16_t *)B, ldb,
+                (float *)C, ldc};
+            return tb.matmul(m, n);
+        }
 #elif defined(__MMA__)
         if (k % 8) {
             return false;
@@ -4054,6 +4150,13 @@ bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64
                 (float *)C, ldc};
             return tb.matmul(m, n);
         }
+        if (Btype == GGML_TYPE_F32) {
+            tinyBLAS<16, __m512, __m512, ggml_fp16_t, float, float> tb{ params, k,
+                (const ggml_fp16_t *)A, lda,
+                (const float *)B, ldb,
+                (float *)C, ldc};
+            return tb.matmul(m, n);
+        }
 #elif (defined(__AVX__) || defined(__AVX2__)) && defined(__F16C__)
         if (Btype == GGML_TYPE_F16) {
             tinyBLAS<8, __m256, __m256, ggml_fp16_t, ggml_fp16_t, float> tb{ params, k,
@@ -4062,6 +4165,13 @@ bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64
                 (float *)C, ldc};
             return tb.matmul(m, n);
         }
+        if (Btype == GGML_TYPE_F32) {
+            tinyBLAS<8, __m256, __m256, ggml_fp16_t, float, float> tb{ params, k,
+                (const ggml_fp16_t *)A, lda,
+                (const float *)B, ldb,
+                (float *)C, ldc};
+            return tb.matmul(m, n);
+        }
 #elif defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER)
         if (n < 8)
             return false;
@@ -4148,6 +4258,17 @@ bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64
             params->ith, params->nth};
         tb.matmul(m, n);
         return true;
+#elif defined(__VXE__) || defined(__VXE2__)
+        if (m % 2 || n % 2) {
+            return false;
+        }
+        tinyBLAS_Q0_VXE<block_q8_0> tb{
+            k, (const block_q8_0 *)A, lda,
+            (const block_q8_0 *)B, ldb,
+            (float *)C, ldc,
+            params->ith, params->nth};
+        tb.matmul(m, n);
+        return true;
 #elif defined(__MMA__)
     //TO-DO: Remove this condition once gemv forwarding is enabled.
         if (n < 8 && n != 4)

ggml/src/ggml-cpu/ops.cpp

@@ -6514,17 +6514,28 @@ static void ggml_compute_forward_conv_transpose_1d_f16_f32(
     const int ith = params->ith;
     const int nth = params->nth;
 
-    const int nk = ne00*ne01*ne02;
-
     GGML_ASSERT(nb00 == sizeof(ggml_fp16_t));
     GGML_ASSERT(nb10 == sizeof(float));
 
     if (ith == 0) {
         memset(params->wdata, 0, params->wsize);
 
+        // permute source data (src1) from (L x Cin) to (Cin x L)
+        {
+            float * const wdata = (float *) params->wdata + 0;
+            float * dst_data = wdata;
+
+            for (int64_t i11 = 0; i11 < ne11; i11++) {
+                const float * const src = (float *)((char *) src1->data + i11*nb11);
+                for (int64_t i10 = 0; i10 < ne10; i10++) {
+                    dst_data[i10*ne11 + i11] = src[i10];
+                }
+            }
+        }
+
         // permute kernel data (src0) from (K x Cout x Cin) to (Cin x K x Cout)
         {
-            ggml_fp16_t * const wdata = (ggml_fp16_t *) params->wdata + 0;
+            ggml_fp16_t * const wdata = (ggml_fp16_t *)((char *) params->wdata + ne10*ne11*sizeof(float));
 
             for (int64_t i02 = 0; i02 < ne02; i02++) {
                 for (int64_t i01 = 0; i01 < ne01; i01++) {
@@ -6537,19 +6548,6 @@ static void ggml_compute_forward_conv_transpose_1d_f16_f32(
             }
         }
 
-        // permute source data (src1) from (L x Cin) to (Cin x L)
-        {
-            ggml_fp16_t * const wdata = (ggml_fp16_t *) params->wdata + nk;
-            ggml_fp16_t * dst_data = wdata;
-
-            for (int64_t i11 = 0; i11 < ne11; i11++) {
-                const float * const src = (float *)((char *) src1->data + i11*nb11);
-                for (int64_t i10 = 0; i10 < ne10; i10++) {
-                    dst_data[i10*ne11 + i11] = GGML_CPU_FP32_TO_FP16(src[i10]);
-                }
-            }
-        }
-
         // need to zero dst since we are accumulating into it
         memset(dst->data, 0, ggml_nbytes(dst));
     }
@@ -6567,19 +6565,19 @@ static void ggml_compute_forward_conv_transpose_1d_f16_f32(
     const int ir0 = dr*ith;
     const int ir1 = MIN(ir0 + dr, nr);
 
-    ggml_fp16_t * const wdata     = (ggml_fp16_t *) params->wdata + 0;
-    ggml_fp16_t * const wdata_src = wdata + nk;
+    float * const wdata_src = (float *) params->wdata + 0;
+    ggml_fp16_t * const wdata_kernel = (ggml_fp16_t *)((char *) params->wdata + ne10*ne11*sizeof(float));
 
     for (int i1 = ir0; i1 < ir1; i1++) {
         float * dst_data = (float *)((char *) dst->data + i1*nb1);
-        ggml_fp16_t * wdata_kernel = wdata + i1*ne02*ne00;
+        ggml_fp16_t * wdata_kernel_row = wdata_kernel + i1*ne02*ne00;
         for (int i10 = 0; i10 < ne10; i10++) {
             const int i1n = i10*ne11;
             for (int i00 = 0; i00 < ne00; i00++) {
                 float v = 0;
-                ggml_vec_dot_f16(ne02, &v, 0,
-                        (ggml_fp16_t *)    wdata_src + i1n, 0,
-                        (ggml_fp16_t *) wdata_kernel + i00*ne02, 0, 1);
+                ggml_vec_dot_f16_f32(ne02, &v, 0,
+                        wdata_kernel_row + i00*ne02, 0,
+                        wdata_src + i1n, 0, 1);
                 dst_data[i10*s0 + i00] += v;
             }
         }
@@ -7172,29 +7170,30 @@ void ggml_compute_forward_im2col_3d(
     }
 }
 
-static void ggml_call_mul_mat(ggml_type type, const ggml_compute_params * params, int64_t m, int64_t n, int64_t k,
+static void ggml_call_mul_mat(ggml_type src0_type, ggml_type src1_type, const ggml_compute_params * params, int64_t m, int64_t n, int64_t k,
                               void * a, void * b, float * c) {
-    const ggml_type_traits * traits = ggml_get_type_traits(type);
+    const ggml_type_traits * traits0 = ggml_get_type_traits(src0_type);
+    const ggml_type_traits * traits1 = ggml_get_type_traits(src1_type);
     struct ggml_tensor src1 = {};
-    src1.type  = type;
+    src1.type  = src1_type;
     src1.ne[0] = k;
     src1.ne[1] = m;
     src1.ne[2] = 1;
     src1.ne[3] = 1;
-    src1.nb[0] = traits->type_size;
-    src1.nb[1] = k * traits->type_size;
+    src1.nb[0] = traits1->type_size;
+    src1.nb[1] = k * traits1->type_size;
     src1.nb[2] = src1.nb[1];
     src1.nb[3] = src1.nb[2];
     src1.data  = a;
 
     struct ggml_tensor src0 = {};
-    src0.type  = type;
+    src0.type  = src0_type;
     src0.ne[0] = k;
     src0.ne[1] = n;
     src0.ne[2] = 1;
     src0.ne[3] = 1;
-    src0.nb[0] = traits->type_size;
-    src0.nb[1] = k * traits->type_size;
+    src0.nb[0] = traits0->type_size;
+    src0.nb[1] = k * traits0->type_size;
     src0.nb[2] = src0.nb[1];
     src0.nb[3] = src0.nb[2];
     src0.data  = b;
@@ -7304,8 +7303,6 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
     GGML_ASSERT(kernel_type == GGML_TYPE_F16 || kernel_type == GGML_TYPE_F32);
     GGML_ASSERT(kernel->type == kernel_type);
 
-    const ggml_type_traits * traits = ggml_get_type_traits(kernel_type);
-
     const int32_t stride_x   = dst->op_params[0];
     const int32_t stride_y   = dst->op_params[1];
     const int32_t pad_x      = dst->op_params[2];
@@ -7331,7 +7328,7 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
     const int64_t knl_n           = knl_w * knl_h * c_in;
     const int64_t patch_total     = dst->ne[3] * dst_w * dst_h;
 
-    const int64_t space_per_patch   = knl_n * traits->type_size + c_out * sizeof(float);
+    const int64_t space_per_patch   = knl_n * sizeof(float) + c_out * sizeof(float);
     const int64_t batch_size        = params->wsize / space_per_patch;
     const int64_t patches_per_batch = batch_size > 8 ? (batch_size / 8) * 8 : batch_size;
     const int64_t batch_n           = (patch_total + patches_per_batch - 1) / patches_per_batch;
@@ -7358,7 +7355,7 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
             const int64_t  src_y       =  p % dst_w;
 
             const float * src_base = (const float *)((const char *)src_data + batch_n * src->nb[3]);
-            char *        dst_row  = (char *) tmp + (p % patches_per_batch) * knl_n * traits->type_size;
+            float *       dst_row  = (float *) tmp + (p % patches_per_batch) * knl_n;
 
             for (int64_t ic = 0; ic < c_in; ++ic) {
                 for (int64_t ky = 0; ky < knl_h; ++ky) {
@@ -7376,12 +7373,7 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
                             src_val               = *src_ptr;
                         }
 
-                        char * element_ptr = dst_row + dst_idx * traits->type_size;
-                        if (kernel_type == GGML_TYPE_F32) {
-                            *(float *) element_ptr = src_val;
-                        } else if (kernel_type == GGML_TYPE_F16) {
-                            *(ggml_fp16_t *) element_ptr = GGML_CPU_FP32_TO_FP16(src_val);
-                        }
+                        dst_row[dst_idx] = src_val;
                     }
                 }
             }
@@ -7389,12 +7381,12 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
 
         ggml_barrier(params->threadpool);
 
-        float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n * traits->type_size);
+        float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n * sizeof(float));
 
         GGML_ASSERT(gemm_output + patch_n * c_out <= (float*)tmp + params->wsize);
 
         // GEMM: patches[patch_n, knl_n] × kernel[knl_n, c_out] = output[patch_n, c_out]
-        ggml_call_mul_mat(kernel_type, params, patch_n, c_out, knl_n, tmp, knl_data, gemm_output);
+        ggml_call_mul_mat(kernel_type, GGML_TYPE_F32, params, patch_n, c_out, knl_n, tmp, knl_data, gemm_output);
 
         ggml_barrier(params->threadpool);
 
@@ -7441,8 +7433,6 @@ static void ggml_compute_forward_conv_3d_impl(const ggml_compute_params * params
     GGML_ASSERT(kernel_type == GGML_TYPE_F16 || kernel_type == GGML_TYPE_F32);
     GGML_ASSERT(kernel->type == kernel_type);
 
-    const ggml_type_traits * traits = ggml_get_type_traits(kernel_type);
-
     const int32_t s0 = dst->op_params[0];
     const int32_t s1 = dst->op_params[1];
     const int32_t s2 = dst->op_params[2];
@@ -7474,7 +7464,7 @@ static void ggml_compute_forward_conv_3d_impl(const ggml_compute_params * params
     const int64_t knl_n_total       = knl_n_per_channel * c;
     const int64_t patch_total       = n * dst_w * dst_h * dst_d;
 
-    const int64_t space_per_patch   = knl_n_total * traits->type_size + oc * sizeof(float);
+    const int64_t space_per_patch   = knl_n_total * sizeof(float) + oc * sizeof(float);
     const int64_t batch_size        = params->wsize / space_per_patch;
     const int64_t patches_per_batch = batch_size > 8 ? (batch_size / 8) * 8 : batch_size;
     const int64_t batch_n           = (patch_total + patches_per_batch - 1) / patches_per_batch;
@@ -7500,7 +7490,7 @@ static void ggml_compute_forward_conv_3d_impl(const ggml_compute_params * params
             const int64_t dst_y      = p_in_depth / dst_w;
             const int64_t dst_x      = p_in_depth % dst_w;
 
-            char * dst_row = (char *) tmp + (p % patches_per_batch) * knl_n_total * traits->type_size;
+            float * dst_row = (float *) tmp + (p % patches_per_batch) * knl_n_total;
 
             for (int64_t ic = 0; ic < c; ++ic) {
                 for (int64_t kz = 0; kz < knl_d; ++kz) {
@@ -7521,12 +7511,7 @@ static void ggml_compute_forward_conv_3d_impl(const ggml_compute_params * params
                                 src_val = *src_ptr;
                             }
 
-                            char * element_ptr = dst_row + dst_idx * traits->type_size;
-                            if (kernel_type == GGML_TYPE_F32) {
-                                *(float *)element_ptr = src_val;
-                            } else if (kernel_type == GGML_TYPE_F16) {
-                                *(ggml_fp16_t *)element_ptr = GGML_CPU_FP32_TO_FP16(src_val);
-                            }
+                            dst_row[dst_idx] = src_val;
                         }
                     }
                 }
@@ -7535,8 +7520,8 @@ static void ggml_compute_forward_conv_3d_impl(const ggml_compute_params * params
 
         ggml_barrier(params->threadpool);
 
-        float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n_total * traits->type_size);
-        ggml_call_mul_mat(kernel_type, params, patch_n_in_batch, oc, knl_n_total, tmp, knl_data, gemm_output);
+        float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n_total * sizeof(float));
+        ggml_call_mul_mat(kernel_type, GGML_TYPE_F32, params, patch_n_in_batch, oc, knl_n_total, tmp, knl_data, gemm_output);
 
         ggml_barrier(params->threadpool);
 
@@ -7588,7 +7573,7 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
     const int ith = params->ith;
     const int nth = params->nth;
 
-    const int nk = ne00*ne01*ne02*ne03;
+    const int nsrc = ne10*ne11*ne12*ne13;
 
     GGML_ASSERT(nb00 == ggml_type_size(src0->type));
     GGML_ASSERT(nb10 == sizeof(float));
@@ -7596,9 +7581,26 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
     if (ith == 0) {
         memset(params->wdata, 0, params->wsize);
 
+        // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh), for all batches
+        {
+            float * const wdata = (float *) params->wdata + 0;
+            for (int i13 = 0; i13 < ne13; i13++) {
+                float * const wdata_b = wdata + i13*ne10*ne11*ne12;
+                for (int i12 = 0; i12 < ne12; i12++) {
+                    for (int i11 = 0; i11 < ne11; i11++) {
+                        const float * const src = (float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11);
+                        float * dst_data = wdata_b + i11*ne10*ne12;
+                        for (int i10 = 0; i10 < ne10; i10++) {
+                            dst_data[i10*ne12 + i12] = src[i10];
+                        }
+                    }
+                }
+            }
+        }
+
         // permute kernel data (src0) from (Kw x Kh x Cout x Cin) to (Cin x Kw x Kh x Cout)
         {
-            kernel_t * const wdata = (kernel_t *) params->wdata + 0;
+            kernel_t * const wdata = (kernel_t *)((char *) params->wdata + nsrc*sizeof(float));
 
             for (int64_t i03 = 0; i03 < ne03; i03++) {
                 for (int64_t i02 = 0; i02 < ne02; i02++) {
@@ -7613,27 +7615,6 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
             }
         }
 
-        // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh), for all batches
-        {
-            kernel_t * const wdata = (kernel_t *) params->wdata + nk;
-            for (int i13 = 0; i13 < ne13; i13++) {
-                kernel_t * const wdata_b = wdata + i13*ne10*ne11*ne12;
-                for (int i12 = 0; i12 < ne12; i12++) {
-                    for (int i11 = 0; i11 < ne11; i11++) {
-                        const float * const src = (float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11);
-                        kernel_t * dst_data = wdata_b + i11*ne10*ne12;
-                        for (int i10 = 0; i10 < ne10; i10++) {
-                            if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
-                                dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]);
-                            } else {
-                                dst_data[i10*ne12 + i12] = src[i10];
-                            }
-                        }
-                    }
-                }
-            }
-        }
-
         memset(dst->data, 0, ggml_nbytes(dst));
     }
     ggml_barrier(params->threadpool);
@@ -7650,14 +7631,14 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
     const int ip0 = dp*ith;
     const int ip1 = MIN(ip0 + dp, np);
 
-    kernel_t * const wdata = (kernel_t *) params->wdata + 0;
-    kernel_t * const wdata_src = wdata + nk;
+    float * const wdata_src = (float *) params->wdata + 0;
+    kernel_t * const wdata_kernel = (kernel_t *)((char *) params->wdata + nsrc*sizeof(float));
 
     for (int i2 = ip0; i2 < ip1; i2++) { // Cout
-        kernel_t * wdata_kernel = wdata + i2*ne01*ne00*ne03;
+        kernel_t * wdata_kernel_c = wdata_kernel + i2*ne01*ne00*ne03;
         for (int i3 = 0; i3 < ne3; i3++) { // batch
             float * dst_data = (float *)((char *) dst->data + i3*nb3 + i2*nb2);
-            kernel_t * wdata_src_b = wdata_src + i3*ne10*ne11*ne12;
+            float * wdata_src_b = wdata_src + i3*ne10*ne11*ne12;
             for (int i11 = 0; i11 < ne11; i11++) {
                 for (int i10 = 0; i10 < ne10; i10++) {
                     const int i1n = i11*ne10*ne12 + i10*ne12;
@@ -7665,13 +7646,13 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
                         for (int i00 = 0; i00 < ne00; i00++) {
                             float v = 0;
                             if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
-                                ggml_vec_dot_f16(ne03, &v, 0,
-                                        wdata_src_b + i1n, 0,
-                                        wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
+                                ggml_vec_dot_f16_f32(ne03, &v, 0,
+                                        wdata_kernel_c + i01*ne00*ne03 + i00*ne03, 0,
+                                        wdata_src_b + i1n, 0, 1);
                             } else {
                                 ggml_vec_dot_f32(ne03, &v, 0,
-                                        wdata_src_b + i1n, 0,
-                                        wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
+                                        wdata_kernel_c + i01*ne00*ne03 + i00*ne03, 0,
+                                        wdata_src_b + i1n, 0, 1);
                             }
                             dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v;
                         }
@@ -9591,6 +9572,7 @@ void ggml_compute_forward_flash_attn_ext(
     switch (dst->op_params[3]) {
         case GGML_PREC_DEFAULT:
         case GGML_PREC_F32:
+        case GGML_PREC_BF16:
             {
                 // uses F32 accumulators
                 ggml_compute_forward_flash_attn_ext_f16(params, dst);

ggml/src/ggml-cpu/spacemit/ime.cpp

@@ -198,7 +198,7 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
                     case GGML_TYPE_Q5_0:
                     case GGML_TYPE_Q5_1:
                     case GGML_TYPE_Q5_K:
-                        //case GGML_TYPE_MXFP4:
+                    case GGML_TYPE_MXFP4:
                         forward_mul_mat(params, op);
                         return true;
                     default:
@@ -1180,7 +1180,7 @@ class tensor_traits_common : public tensor_traits_base {
         const int64_t DK = nek0;
         const int64_t DV = nev0;
 
-        const bool supported_prec  = (dst->op_params[3] == GGML_PREC_F32 || dst->op_params[3] == GGML_PREC_DEFAULT);
+        const bool supported_prec  = (dst->op_params[3] == GGML_PREC_F32 || dst->op_params[3] == GGML_PREC_BF16 || dst->op_params[3] == GGML_PREC_DEFAULT);
         const bool supported_types = (q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16);
         const bool supported_shape = (DK > 0 && DK <= 128 && DV > 0 && DV <= 128);
         const bool supported_vlen  = (__riscv_vlenb() == 128);
@@ -1375,10 +1375,10 @@ static const ggml::cpu::tensor_traits * ggml_riscv64_spacemit_get_optimal_repack
         case GGML_TYPE_MXFP4:
             {
 #if defined(RISCV64_SPACEMIT_IME2)
-                // TODO
-                // if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
-                //     return &ggml::cpu::riscv64_spacemit::mxfp4_32x32_q8_0;
-                // }
+                if (ggml_n_dims(cur) == 2 && cur->ne[0] % 32 == 0 && cur->ne[1] % 32 == 0 &&
+                    (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::mxfp4_32x32_q8_0;
+                }
 #endif
             }
             break;
@@ -1613,6 +1613,9 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
                     if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) {
                         return false;
                     }
+                    if (op->src[0]->type == GGML_TYPE_MXFP4 && !ggml_is_contiguous(op->src[1])) {
+                        return false;
+                    }
                     if (op->src[1]->type == GGML_TYPE_F32) {
                         return true;
                     }

ggml/src/ggml-cpu/spacemit/ime2_kernels.cpp

@@ -3754,7 +3754,7 @@ void gemm_kernel_i8mxfp4_m1(size_t          blk_len,
     using blk_type           = nrow_block_mxfp4<NB_COLS>;
 
     GGML_ASSERT(blk_len == K_TILE);
-    GGML_ASSERT(count_m == 1);
+    GGML_ASSERT(count_m >= 1);
     GGML_UNUSED(quant_b_zp);
 
     const size_t a_blk_stride  = q8_blk_size(blk_len, true);
@@ -3923,7 +3923,7 @@ void gemm_kernel_i8mxfp4_m1(size_t          blk_len,
                 :
                 : [BCK] "r"(cnt), [NBLKS] "r"(nblks), [pA] "r"(quant_a_ptr), [pB] "r"(QuantBDataPtr), [pC] "r"(CPtr)
                 : "cc", "memory", "t0", "t1", "t2", "t3", "t4", "f0", "s2", "s3", "s4", "s5", "s6", "s7", "v0", "v1",
-                  "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v16", "v17", "v18", "v19",
+                  "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19",
                   "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31");
         }
     }
@@ -3943,7 +3943,7 @@ void gemm_kernel_i8mxfp4_m4(size_t          blk_len,
     using blk_type           = nrow_block_mxfp4<NB_COLS>;
 
     GGML_ASSERT(blk_len == K_TILE);
-    GGML_ASSERT(count_m == 4);
+    GGML_ASSERT(count_m >= 4);
     GGML_UNUSED(quant_b_zp);
 
     const size_t a_blk_stride  = q8_blk_size(blk_len, true);
@@ -5687,20 +5687,10 @@ size_t gemm_kernel_i8mxfp4(size_t          blk_len,
                            size_t          k_blks,
                            size_t          ldc) {
     if (count_m >= 4) {
-#if 1
-        gemm_kernel_i8mxfp4_mrow_ref<4, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
-                                            k_blks, ldc);
-#else
         gemm_kernel_i8mxfp4_m4(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
-#endif
         return 4;
     } else {
-#if 1
-        gemm_kernel_i8mxfp4_mrow_ref<1, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
-                                            k_blks, ldc);
-#else
         gemm_kernel_i8mxfp4_m1(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
-#endif
         return 1;
     }
 }

ggml/src/ggml-cpu/vec.cpp

@@ -377,6 +377,85 @@ void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * G
     *s = sumf;
 }
 
+void ggml_vec_dot_f16_f32(int n, float * GGML_RESTRICT s, size_t bs, const ggml_fp16_t * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc) {
+    assert(nrc == 1);
+    GGML_UNUSED(nrc);
+    GGML_UNUSED(bx);
+    GGML_UNUSED(by);
+    GGML_UNUSED(bs);
+
+    ggml_float sumf = 0.0;
+    int i = 0;
+
+    // on AVX512 and AVX with F16C, x is widened to F32 in-register and accumulated in F32 over 4 independent sums
+#if defined(__AVX512F__)
+    __m512 acc0 = _mm512_setzero_ps();
+    __m512 acc1 = _mm512_setzero_ps();
+    __m512 acc2 = _mm512_setzero_ps();
+    __m512 acc3 = _mm512_setzero_ps();
+    for (; i + 64 <= n; i += 64) {
+        acc0 = _mm512_fmadd_ps(_mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x + i +  0))), _mm512_loadu_ps(y + i +  0), acc0);
+        acc1 = _mm512_fmadd_ps(_mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x + i + 16))), _mm512_loadu_ps(y + i + 16), acc1);
+        acc2 = _mm512_fmadd_ps(_mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x + i + 32))), _mm512_loadu_ps(y + i + 32), acc2);
+        acc3 = _mm512_fmadd_ps(_mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x + i + 48))), _mm512_loadu_ps(y + i + 48), acc3);
+    }
+    for (; i + 16 <= n; i += 16) {
+        acc0 = _mm512_fmadd_ps(_mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x + i))), _mm512_loadu_ps(y + i), acc0);
+    }
+    sumf = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3)));
+#elif defined(__AVX__) && defined(__F16C__)
+    __m256 acc0 = _mm256_setzero_ps();
+    __m256 acc1 = _mm256_setzero_ps();
+    __m256 acc2 = _mm256_setzero_ps();
+    __m256 acc3 = _mm256_setzero_ps();
+    for (; i + 32 <= n; i += 32) {
+        acc0 = GGML_F32x8_FMA(acc0, _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i +  0))), _mm256_loadu_ps(y + i +  0));
+        acc1 = GGML_F32x8_FMA(acc1, _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i +  8))), _mm256_loadu_ps(y + i +  8));
+        acc2 = GGML_F32x8_FMA(acc2, _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i + 16))), _mm256_loadu_ps(y + i + 16));
+        acc3 = GGML_F32x8_FMA(acc3, _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i + 24))), _mm256_loadu_ps(y + i + 24));
+    }
+    for (; i + 8 <= n; i += 8) {
+        acc0 = GGML_F32x8_FMA(acc0, _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i))), _mm256_loadu_ps(y + i));
+    }
+    const __m256 acc = _mm256_add_ps(_mm256_add_ps(acc0, acc1), _mm256_add_ps(acc2, acc3));
+    __m128 r = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps(acc, 1));
+    r = _mm_add_ps(r, _mm_movehl_ps(r, r));
+    r = _mm_add_ss(r, _mm_movehdup_ps(r));
+    sumf = _mm_cvtss_f32(r);
+#elif defined(GGML_SIMD)
+    // convert chunks of x to F32 and reuse the F32 SIMD dot
+    // TODO: convert F16 to F32 in-register and fuse the dot (needs a new SIMD macro); also avoid the per-chunk call overhead for tiny k
+    float tx[64];
+    for (; i + 64 <= n; i += 64) {
+        const ggml_fp16_t * GGML_RESTRICT xr = x + i;
+        float * GGML_RESTRICT tr = tx;
+        for (int j = 0; j < 64; ++j) {
+            tr[j] = GGML_CPU_FP16_TO_FP32(xr[j]);
+        }
+        float v;
+        ggml_vec_dot_f32(64, &v, 0, tx, 0, y + i, 0, 1);
+        sumf += v;
+    }
+    if (i < n) {
+        const int nr = n - i;
+        const ggml_fp16_t * GGML_RESTRICT xr = x + i;
+        float * GGML_RESTRICT tr = tx;
+        for (int j = 0; j < nr; ++j) {
+            tr[j] = GGML_CPU_FP16_TO_FP32(xr[j]);
+        }
+        float v;
+        ggml_vec_dot_f32(nr, &v, 0, tx, 0, y + i, 0, 1);
+        sumf += v;
+        i = n;
+    }
+#endif
+    for (; i < n; ++i) {
+        sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*y[i]);
+    }
+
+    *s = sumf;
+}
+
 void ggml_vec_silu_f32(const int n, float * y, const float * x) {
     int i = 0;
 #if defined(__AVX512F__) && defined(__AVX512DQ__)

ggml/src/ggml-cpu/vec.h

@@ -70,6 +70,7 @@ extern ggml_fp16_t ggml_table_gelu_quick_f16[1 << 16];
 void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc);
 void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t * GGML_RESTRICT x, size_t bx, ggml_bf16_t * GGML_RESTRICT y, size_t by, int nrc);
 void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * GGML_RESTRICT x, size_t bx, ggml_fp16_t * GGML_RESTRICT y, size_t by, int nrc);
+void ggml_vec_dot_f16_f32(int n, float * GGML_RESTRICT s, size_t bs, const ggml_fp16_t * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc);
 
 void ggml_vec_silu_f32(const int n, float * y, const float * x);
 ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean); //it will also center y ( y = y - mean )

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -5978,7 +5978,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                 if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16 && !ggml_cuda_op_mul_mat_use_fwht(op)) {
                     return false;
                 }
-                if (op->op == GGML_OP_MUL_MAT_ID && ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) {
+                if (op->op == GGML_OP_MUL_MAT_ID && (ggml_get_op_params_i32(op, 3) == GGML_PREC_BF16 || ggml_get_op_params_i32(op, 3) == GGML_PREC_F32)) {
                     return false;
                 }
 #ifdef GGML_USE_MUSA
@@ -6252,9 +6252,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
 
             if (op->src[3]->ne[0] == 1) {
                 // Mamba2
-                // (kernel only supports (d_state == 96 || d_state == 128 || d_state == 256) && d_head % 16 == 0)
+                // (kernel only supports (d_state == 64 || d_state == 96 || d_state == 128 || d_state == 256) && d_head % 16 == 0)
                 const int64_t d_state = op->src[0]->ne[0];
-                return (d_state == 96 || d_state == 128 || d_state == 256) && op->src[0]->ne[1] % 16 == 0;
+                return (d_state == 64 || d_state == 96 || d_state == 128 || d_state == 256) && op->src[0]->ne[1] % 16 == 0;
             } else {
                 if (K > 1) {
                     return false;

ggml/src/ggml-cuda/ssm-scan.cu

@@ -256,7 +256,17 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
     // NOTE: if you change conditions here, be sure to update the corresponding supports_op condition!
     if (src3_nb1 == sizeof(float)) {
         // Mamba-2
-        if (d_state == 96) {
+        if (d_state == 64) { // Zamba2
+            constexpr int threads   = 64;
+            constexpr int num_warps = threads/WARP_SIZE;
+
+            const dim3 blocks((n_head * head_dim + (num_warps - 1)) / num_warps, n_seq, 1);
+            const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, 0, stream);
+            ggml_cuda_kernel_launch(ssm_scan_f32_group<64/WARP_SIZE, 64>, launch_params,
+                    src0, src1, src2, src3, src4, src5, src6, dst,
+                    src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
+                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_base, s_slot_bytes, n_head, head_dim, n_group, n_tok, K);
+        } else if (d_state == 96) {
             constexpr int threads   = 96;
             constexpr int num_warps = threads/WARP_SIZE;
 
@@ -287,7 +297,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
                     src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
                     src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_base, s_slot_bytes, n_head, head_dim, n_group, n_tok, K);
         } else {
-            GGML_ABORT("doesn't support d_state!=(96, 128 or 256).");
+            GGML_ABORT("doesn't support d_state!=(64, 96, 128 or 256).");
         }
     } else {
         // Mamba-1

ggml/src/ggml-cuda/top-k.cu

@@ -88,18 +88,18 @@ static __device__ __forceinline__ uint32_t top_k_float_to_ordered(float value) {
     return bits ^ mask;
 }
 
-struct top_k_radix_state {
+// one state per cache line, so the gather atomics of different rows do not compete
+struct alignas(128) top_k_radix_state {
     uint32_t prefix;
     uint32_t prefix_mask;
     int rank;
     int greater_count;
-    int equal_count;
 };
 
 static __global__ void top_k_radix_init(top_k_radix_state * states, int nrows, int k) {
     const int row = blockIdx.x * blockDim.x + threadIdx.x;
     if (row < nrows) {
-        states[row] = {0, 0, k, 0, 0};
+        states[row] = {0, 0, k, 0};
     }
 }
 
@@ -109,6 +109,7 @@ static __global__ void top_k_radix_histogram(
         const top_k_radix_state * __restrict__ states,
         int * __restrict__ block_histograms,
         int ncols,
+        int block_cols,
         int blocks_per_row,
         int shift) {
     constexpr int NBINS = 1 << RADIX_BITS;
@@ -117,15 +118,15 @@ static __global__ void top_k_radix_histogram(
     const int row_block = blockIdx.x % blocks_per_row;
     const int tid = threadIdx.x;
     const float * row_src = src + (size_t) row * ncols;
+    const int64_t col_begin = (int64_t) row_block * block_cols;
+    const int64_t col_end   = min(col_begin + block_cols, (int64_t) ncols);
     __shared__ int histogram[NBINS];
 
     histogram[tid] = 0;
     __syncthreads();
 
     const top_k_radix_state state = states[row];
-    for (int64_t col = row_block * BLOCK_SIZE + tid;
-         col < ncols;
-         col += blocks_per_row * BLOCK_SIZE) {
+    for (int64_t col = col_begin + tid; col < col_end; col += BLOCK_SIZE) {
         const uint32_t key = top_k_float_to_ordered(row_src[col]);
         if ((key & state.prefix_mask) == state.prefix) {
             atomicAdd(&histogram[(key >> shift) & (NBINS - 1)], 1);
@@ -170,43 +171,91 @@ static __global__ void top_k_radix_select(
     }
 }
 
-static __global__ void top_k_radix_reset_counters(top_k_radix_state * states, int nrows) {
-    const int row = blockIdx.x * blockDim.x + threadIdx.x;
-    if (row < nrows) {
-        states[row].greater_count = 0;
-        states[row].equal_count = 0;
-    }
-}
-
-template<int BLOCK_SIZE>
+template<int BLOCK_SIZE, int RADIX_BITS>
 static __global__ void top_k_radix_gather(
         const float * __restrict__ src,
         int * __restrict__ dst,
+        const int * __restrict__ block_histograms,
         top_k_radix_state * __restrict__ states,
         int ncols,
+        int block_cols,
         int k,
         int blocks_per_row) {
+    constexpr int NBINS     = 1 << RADIX_BITS;
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = BLOCK_SIZE / warp_size;
+    using lane_mask_t = std::conditional_t<warp_size == 64, uint64_t, uint32_t>;
+    // HIP requires a 64-bit mask also for a 32-wide warp
+    constexpr uint64_t ballot_mask = ~(lane_mask_t) 0;
+    const auto popc = [](const lane_mask_t mask) {
+        if constexpr (warp_size == 64) {
+            return __popcll(mask);
+        } else {
+            return __popc(mask);
+        }
+    };
+    __shared__ int warp_ties[nwarps];
+
     const int row = blockIdx.x / blocks_per_row;
     const int row_block = blockIdx.x % blocks_per_row;
-    const int tid = threadIdx.x;
+    const int lane = threadIdx.x % warp_size;
+    const int warp = threadIdx.x / warp_size;
     const float * row_src = src + (size_t) row * ncols;
     int * row_dst = dst + (size_t) row * k;
     top_k_radix_state * state = &states[row];
+    const uint32_t threshold = state->prefix;
+    const int rank = state->rank;
 
-    for (int64_t col = row_block * BLOCK_SIZE + tid;
-         col < ncols;
-         col += blocks_per_row * BLOCK_SIZE) {
+    // each warp reads a contiguous part of the contiguous range of the block
+    const int     warp_cols  = block_cols / nwarps;
+    const int64_t warp_begin = (int64_t) row_block * block_cols + warp * warp_cols;
+    const int64_t warp_end   = min(warp_begin + warp_cols, (int64_t) ncols);
+
+    int nties = 0;
+    for (int64_t col = warp_begin + lane; col < warp_end; col += warp_size) {
         const uint32_t key = top_k_float_to_ordered(row_src[col]);
-        if (key > state->prefix) {
+        if (key > threshold) {
             const int pos = atomicAdd(&state->greater_count, 1);
             row_dst[pos] = col;
-        } else if (key == state->prefix) {
-            const int pos = atomicAdd(&state->equal_count, 1);
-            if (pos < state->rank) {
-                row_dst[k - state->rank + pos] = col;
-            }
+        } else if (key == threshold) {
+            nties++;
         }
     }
+
+    // the ties fill the last rank slots in column order, the lowest columns win.
+    // the histograms of the last pass count the ties of each block in the bin of the threshold.
+    const int * block_ties = block_histograms + (size_t) row * blocks_per_row * NBINS + (threshold & (NBINS - 1));
+    int pos = 0;
+    for (int b = lane; b < row_block; b += warp_size) {
+        pos += block_ties[b * NBINS];
+    }
+    pos = warp_reduce_sum<warp_size>(pos);
+    if (pos >= rank || block_ties[row_block * NBINS] == 0) {
+        return; // the same result in all threads of the block
+    }
+
+    nties = warp_reduce_sum<warp_size>(nties);
+    if (lane == 0) {
+        warp_ties[warp] = nties;
+    }
+    __syncthreads();
+    for (int w = 0; w < warp; ++w) {
+        pos += warp_ties[w];
+    }
+
+    // pos is the same in all lanes of the warp, so the lanes leave the loop together
+    const int pos_end = min(pos + nties, rank);
+    for (int64_t tile = warp_begin; tile < warp_end && pos < pos_end; tile += warp_size) {
+        const int64_t col    = tile + lane;
+        const bool    is_tie = col < warp_end && top_k_float_to_ordered(row_src[col]) == threshold;
+
+        const lane_mask_t ties    = (lane_mask_t) __ballot_sync(ballot_mask, is_tie);
+        const int         tie_pos = pos + popc(ties & (((lane_mask_t) 1 << lane) - 1));
+        if (is_tie && tie_pos < pos_end) {
+            row_dst[k - rank + tie_pos] = col;
+        }
+        pos += popc(ties);
+    }
 }
 
 static void top_k_radix_cuda(
@@ -216,6 +265,8 @@ static void top_k_radix_cuda(
     constexpr int RADIX_BITS = 8;
     constexpr int NBINS = 1 << RADIX_BITS;
     const int blocks_per_row = (int) std::min<int64_t>(((int64_t) ncols + 1023) / 1024, 64);
+    // each block of a row reads a contiguous range of whole tiles
+    const int block_cols = (int) (((int64_t) ncols + blocks_per_row * BLOCK_SIZE - 1) / (blocks_per_row * BLOCK_SIZE) * BLOCK_SIZE);
 
     // chunk the rows to bound the histogram memory to 64 MB
     const int64_t chunk_nrows = ggml_cuda_chunk_nrows((size_t) blocks_per_row * NBINS * sizeof(int), nrows);
@@ -234,16 +285,14 @@ static void top_k_radix_cuda(
         for (int shift = 32 - RADIX_BITS; shift >= 0; shift -= RADIX_BITS) {
             top_k_radix_histogram<BLOCK_SIZE, RADIX_BITS>
                 <<<row_grid, BLOCK_SIZE, 0, stream>>>(
-                    src, states, histograms, ncols, blocks_per_row, shift);
+                    src, states, histograms, ncols, block_cols, blocks_per_row, shift);
             top_k_radix_select<BLOCK_SIZE, RADIX_BITS>
                 <<<iter_nrows, BLOCK_SIZE, 0, stream>>>(histograms, states, blocks_per_row, shift);
         }
 
-        top_k_radix_reset_counters
-            <<<(iter_nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, iter_nrows);
-        top_k_radix_gather<BLOCK_SIZE>
+        top_k_radix_gather<BLOCK_SIZE, RADIX_BITS>
             <<<row_grid, BLOCK_SIZE, 0, stream>>>(
-                src, dst, states, ncols, k, blocks_per_row);
+                src, dst, histograms, states, ncols, block_cols, k, blocks_per_row);
 
         src += (size_t) ncols * iter_nrows;
         dst += (size_t) k     * iter_nrows;

ggml/src/ggml-et/ggml-et-ops.cpp

@@ -1523,7 +1523,7 @@ bool ggml_et_op_flash_attn_ext(ggml_backend_et_device_context * dev_ctx, const g
     }
 
     const ggml_prec prec = ggml_flash_attn_ext_get_prec(node);
-    if (prec != GGML_PREC_F32 && prec != GGML_PREC_DEFAULT) {
+    if (prec != GGML_PREC_F32 && prec != GGML_PREC_BF16 && prec != GGML_PREC_DEFAULT) {
         GGML_LOG_ERROR("ET: FLASH_ATTN_EXT baseline kernel only supports F32 precision\n");
         return false;
     }

ggml/src/ggml-et/ggml-et.cpp

@@ -1299,7 +1299,7 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm
                 const bool me_eligible = op->src[1]->type == GGML_TYPE_F16 && op->src[2]->type == GGML_TYPE_F16 &&
                                          (op->src[0]->ne[0] % 32) == 0;
 
-                supported = me_eligible && mask_ok && (prec == GGML_PREC_F32 || prec == GGML_PREC_DEFAULT) &&
+                supported = me_eligible && mask_ok && (prec == GGML_PREC_F32 || prec == GGML_PREC_BF16 || prec == GGML_PREC_DEFAULT) &&
                             max_bias == 0.0f && logit_softcap == 0.0f && op->src[0]->nb[0] == sizeof(float) &&
                             op->src[1]->nb[0] == k_elem && op->src[2]->nb[0] == v_elem && op->nb[0] == sizeof(float) &&
                             op->src[0]->ne[0] == op->src[1]->ne[0] &&  // dk matches

ggml/src/ggml-hexagon/ggml-hexagon.cpp

@@ -242,10 +242,11 @@ static void ggml_hexagon_dump_batch_prof(const char * sess_name, const htp_opbat
 
     char evt_str[256] = "----";
     if (opt_profile == 3) {
-        snprintf(evt_str, sizeof(evt_str), "evt-cnt %u,%u,%u,%u,%u,%u,%u,%u,%u,%u,%u",
+        static_assert(HTP_OP_MAX_NTHREADS == 10, "current implementation assumes 10 trace counters");
+        snprintf(evt_str, sizeof(evt_str), "evt-cnt %u,%u,%u,%u,%u,%u,%u,%u,%u,%u",
                 rsp.n_traces[0], rsp.n_traces[1], rsp.n_traces[2], rsp.n_traces[3],
                 rsp.n_traces[4], rsp.n_traces[5], rsp.n_traces[6], rsp.n_traces[7],
-                rsp.n_traces[8], rsp.n_traces[9], rsp.n_traces[10]);
+                rsp.n_traces[8], rsp.n_traces[9]);
     }
 
     GGML_LOG_DEBUG("ggml-hex: %s profile-op OPBATCH|----|n-ops %u|%s|----|----|usec %u cycles %llu start %llu mhz %.1f\n",
@@ -253,16 +254,11 @@ static void ggml_hexagon_dump_batch_prof(const char * sess_name, const htp_opbat
 }
 
 static void ggml_hexagon_dump_trace_events(const char * sess_name, const htp_opbatch_rsp & rsp,
-                                           const htp_trace_desc * trace_events, uint32_t n_traces) {
+                                           const htp_trace_desc * trace_events, uint32_t n_threads, uint32_t n_traces) {
     if (opt_profile == 3 && trace_events) {
-        uint32_t valid_cnt[HTP_MAX_NTHREADS + 1] = {0};
-        for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
-            uint32_t count = rsp.n_traces[t];
-            valid_cnt[t] = count > n_traces ? n_traces : count;
-        }
-
-        for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
-            for (uint32_t idx = 0; idx < valid_cnt[t]; idx++) {
+        for (uint32_t t = 0; t < n_threads; t++) {
+            uint32_t count = std::min(rsp.n_traces[t], n_traces);
+            for (uint32_t idx = 0; idx < count; idx++) {
                 const auto & e = trace_events[t * n_traces + idx];
                 bool is_stop = (e.info & 0x8000) != 0;
                 uint16_t info = e.info & 0x7FFF;
@@ -393,7 +389,7 @@ struct ggml_hexagon_session {
 
 const char * htp_session_name(const struct ggml_hexagon_session * sess) { return sess->c_name(); }
 size_t       htp_session_vtcm_size(const struct ggml_hexagon_session * sess) { return sess->vtcm_size; }
-uint32_t     htp_session_n_threads(const struct ggml_hexagon_session * sess) { return sess->n_threads; }
+uint32_t     htp_session_n_hvx(const struct ggml_hexagon_session * sess) { return sess->n_hvx; }
 uint32_t     htp_session_n_hmx(const struct ggml_hexagon_session * sess) { return sess->n_hmx; }
 uint32_t     htp_session_mdev_count(const struct ggml_hexagon_session * sess) { return sess->mdev.count; }
 
@@ -1547,7 +1543,7 @@ struct ggml_hexagon_opqueue {
 
         size_t tr_size = 0;
         if (opt_profile == 3) {
-            tr_size = (HTP_MAX_NTHREADS + 1) * opt_optrace * sizeof(htp_trace_desc);
+            tr_size = sess->n_threads * opt_optrace * sizeof(htp_trace_desc);
         }
 
         shm_blk_size = sizeof(htp_buf_desc)  * n_bufs    +
@@ -1602,7 +1598,7 @@ struct ggml_hexagon_opqueue {
         size_t tr_size = 0;
         if (opt_profile == 3) {
             req.n_traces = opt_optrace;
-            tr_size = (HTP_MAX_NTHREADS + 1) * req.n_traces * sizeof(htp_trace_desc);
+            tr_size = shm_buf->sess->n_threads * req.n_traces * sizeof(htp_trace_desc);
         } else {
             req.n_traces = 0;
         }
@@ -1657,7 +1653,7 @@ struct ggml_hexagon_opqueue {
         uint32_t n_traces = 0;
         if (opt_profile == 3) {
             n_traces = opt_optrace;
-            tr_size = (HTP_MAX_NTHREADS + 1) * n_traces * sizeof(htp_trace_desc);
+            tr_size = shm_buf->sess->n_threads * n_traces * sizeof(htp_trace_desc);
         }
 
         const size_t m_size = b_size + t_size + o_size + p_size + tr_size;
@@ -1686,7 +1682,7 @@ struct ggml_hexagon_opqueue {
                 ggml_hexagon_dump_op_prof(shm_buf->sess->c_name(), ops[i], pd[i]);
             }
 
-            ggml_hexagon_dump_trace_events(shm_buf->sess->c_name(), rsp, trace_events, n_traces);
+            ggml_hexagon_dump_trace_events(shm_buf->sess->c_name(), rsp, trace_events, shm_buf->sess->n_threads, n_traces);
         }
     }
 };
@@ -2121,34 +2117,80 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
         }
     }
 
-    char session_uri[256];
-    {
-        char htp_uri[256];
-        snprintf(htp_uri, sizeof(htp_uri), "file:///libggml-htp-v%u.so?htp_iface_skel_handle_invoke&_modver=1.0", opt_arch);
+    auto open_session = [&](const std::string & prefix, std::string & uri_out) -> int {
+        const std::string htp_uri = prefix + "libggml-htp-v" + std::to_string(opt_arch) +
+                                    ".so?htp_iface_skel_handle_invoke&_modver=1.0";
+
+        std::vector<char> session_uri(htp_uri.size() + 256);
 
         struct remote_rpc_get_uri u = {};
         u.session_id      = this->session_id;
         u.domain_name     = const_cast<char *>(dom_name.c_str());
         u.domain_name_len = dom_name.size();
-        u.module_uri      = const_cast<char *>(htp_uri);
-        u.module_uri_len  = strlen(htp_uri);
-        u.uri             = session_uri;
-        u.uri_len         = sizeof(session_uri);
+        u.module_uri      = const_cast<char *>(htp_uri.c_str());
+        u.module_uri_len  = htp_uri.size();
+        u.uri             = session_uri.data();
+        u.uri_len         = session_uri.size();
 
         int err = remote_session_control(FASTRPC_GET_URI, (void *) &u, sizeof(u));
         if (err != AEE_SUCCESS) {
-            snprintf(session_uri, sizeof(session_uri), "%s&_dom=%s&_session=%u",
-                     htp_uri, dom_name.c_str(), this->session_id);
+            snprintf(session_uri.data(), session_uri.size(), "%s&_dom=%s&_session=%u",
+                     htp_uri.c_str(), dom_name.c_str(), this->session_id);
 
-            GGML_LOG_WARN("ggml-hex: %s failed to get URI (physical %d, virtual %d) : error 0x%x. Falling back to single session URI: %s\n",
-                          this->c_name(), phys_idx, virt_idx, err, session_uri);
+            HEX_VERBOSE("ggml-hex: %s failed to get URI (physical %d, virtual %d) : error 0x%x. Falling back to single session URI: %s\n",
+                        this->c_name(), phys_idx, virt_idx, err, session_uri.data());
+        }
+
+        uri_out = session_uri.data();
+        return htp_iface_open(session_uri.data(), &this->handle);
+    };
+
+    const std::string lib_dir = htp_lib_dir();
+
+    std::string session_uri;
+    int err = AEE_EFAILED;
+
+    // 1. Try opening session with absolute path in URI
+    if (!lib_dir.empty()) {
+        std::string abs_prefix = "file:///" + lib_dir;
+        if (abs_prefix.back() != '/' && abs_prefix.back() != '\\') {
+            abs_prefix += "/";
+        }
+
+        HEX_VERBOSE("ggml-hex: %s trying absolute path URI: %s\n", this->c_name(), abs_prefix.c_str());
+        err = open_session(abs_prefix, session_uri);
+        if (err != AEE_SUCCESS) {
+            HEX_VERBOSE("ggml-hex: %s failed to open session with absolute path : uri %s error 0x%x\n",
+                        this->c_name(), session_uri.c_str(), err);
         }
     }
 
-    // Open session
-    int err = htp_iface_open(session_uri, &this->handle);
+    // 2. Fallback to DSP search path via FASTRPC_MANAGE_DSP_LIBRARY_PATH
+    if (err != AEE_SUCCESS && !lib_dir.empty()) {
+        int sp_err = htpdrv_set_dsp_search_path(this->domain_id, lib_dir, this->c_name());
+        if (sp_err == AEE_SUCCESS) {
+            err = open_session("file:///", session_uri);
+            if (err != AEE_SUCCESS) {
+                HEX_VERBOSE("ggml-hex: %s failed to open session with search path : uri %s error 0x%x\n",
+                            this->c_name(), session_uri.c_str(), err);
+            }
+        }
+    }
+
+    // 3. Fallback to process environment variables
+    if (err != AEE_SUCCESS && !lib_dir.empty()) {
+        HEX_VERBOSE("ggml-hex: %s falling back to process env for DSP search path\n", this->c_name());
+        htpdrv_set_dsp_env_path(lib_dir, this->c_name());
+        err = open_session("file:///", session_uri);
+    }
+
+    // Bare name attempt if lib_dir is empty and nothing was tried yet
+    if (err != AEE_SUCCESS && lib_dir.empty()) {
+        err = open_session("file:///", session_uri);
+    }
+
     if (err != AEE_SUCCESS) {
-        GGML_LOG_ERROR("ggml-hex: %s failed to open session : uri %s error 0x%x\n", this->c_name(), session_uri, err);
+        GGML_LOG_ERROR("ggml-hex: %s failed to open session : uri %s error 0x%x\n", this->c_name(), session_uri.c_str(), err);
         throw std::runtime_error("ggml-hex: failed to open session (see log for details)");
     }
 
@@ -2162,11 +2204,11 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
         unsigned long long hw_vtcm_size = 0;
         int hw_err = htp_iface_hwinfo(this->handle, &hw_n_threads, &hw_n_hvx, &hw_n_hmx, &hw_vtcm_size);
         if (hw_err == 0) {
-            const uint32_t max_n_threads = (std::min)((uint32_t) HTP_MAX_NTHREADS, (uint32_t) hw_n_threads);
-            this->n_threads = opt_nhvx > 0 ? (uint32_t) (std::min)(opt_nhvx, (size_t) max_n_threads) : max_n_threads;
-            this->n_hvx     = this->n_threads;
-            this->n_hmx     = (opt_nhmx != 0) ? (uint32_t)hw_n_hmx : 0;
-            this->vtcm_size = (uint64_t)hw_vtcm_size;
+            this->n_threads          = (std::min)((uint32_t) HTP_OP_MAX_NTHREADS, (uint32_t) hw_n_threads);
+            const uint32_t max_n_hvx = (std::min)({this->n_threads - 1, (uint32_t) hw_n_hvx, (uint32_t) HTP_OP_MAX_NHVX});
+            this->n_hvx              = opt_nhvx > 0 ? (uint32_t) (std::min)(opt_nhvx, (size_t) max_n_hvx) : max_n_hvx;
+            this->n_hmx              = (opt_nhmx != 0) ? (uint32_t)hw_n_hmx : 0;
+            this->vtcm_size          = (uint64_t)hw_vtcm_size;
             if (opt_vtcm_size > 0) {
                 if (opt_vtcm_size > hw_vtcm_size) {
                     GGML_LOG_WARN("ggml-hex: %s VTCM override %zu exceeds HW size %llu\n", this->c_name(), opt_vtcm_size, hw_vtcm_size);
@@ -2178,11 +2220,11 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
                           (unsigned long long)(this->vtcm_size / (1024 * 1024)));
         } else {
             GGML_LOG_WARN("ggml-hex: %s failed to query hwinfo (0x%x), using defaults\n", this->c_name(), hw_err);
-            const uint32_t default_n_threads = (std::min)(8u, (uint32_t) HTP_MAX_NTHREADS);
-            this->n_threads = opt_nhvx > 0 ? (uint32_t) (std::min)(opt_nhvx, (size_t) HTP_MAX_NTHREADS) : default_n_threads;
-            this->n_hvx     = this->n_threads;
-            this->n_hmx     = (opt_nhmx != 0) ? 1 : 0;
-            this->vtcm_size = opt_vtcm_size > 0 ? opt_vtcm_size : 8 * 1024 * 1024;
+            this->n_threads          = (std::min)(8u, (uint32_t) HTP_OP_MAX_NTHREADS);
+            const uint32_t max_n_hvx = (std::min)(this->n_threads - 1, (uint32_t) HTP_OP_MAX_NHVX);
+            this->n_hvx              = opt_nhvx > 0 ? (uint32_t) (std::min)(opt_nhvx, (size_t) max_n_hvx) : (std::min)(4u, max_n_hvx);
+            this->n_hmx              = (opt_nhmx != 0) ? 1 : 0;
+            this->vtcm_size          = opt_vtcm_size > 0 ? opt_vtcm_size : 8 * 1024 * 1024;
         }
     }
 
@@ -2198,7 +2240,7 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
     }
 
     GGML_LOG_INFO("ggml-hex: %s new session : session-id %d domain-id %d uri %s handle 0x%lx\n", this->c_name(),
-                  this->session_id, this->domain_id, session_uri, (unsigned long) this->handle);
+                  this->session_id, this->domain_id, session_uri.c_str(), (unsigned long) this->handle);
 
     const size_t req_q_size = (sizeof(htp_opbatch_req) * opt_opqueue * 2) + 1024;
     const size_t rsp_q_size = (sizeof(htp_opbatch_rsp) * opt_opqueue * 2) + 1024;
@@ -2251,7 +2293,7 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
     this->op_batch = new ggml_hexagon_opbatch(this, opt_opbatch, this->max_vmem);
 
     // Start dspqueue/opbatch processing
-    err = htp_iface_start(this->handle, this->session_id, this->queue_id, this->n_threads, opt_nhmx, opt_vtcm_size, this->max_vmem,
+    err = htp_iface_start(this->handle, this->session_id, this->queue_id, this->n_hvx, opt_nhmx, opt_vtcm_size, this->max_vmem,
                           (uint32_t) opt_power_mode);
     if (err != 0) {
         GGML_LOG_ERROR("ggml-hex: %s failed to start session: 0x%08x\n", this->c_name(), (unsigned) err);
@@ -3353,8 +3395,8 @@ static void ggml_hexagon_discover_devices() {
             discovery_supported = true;
             const int n_domains = std::min(domain_info.sys.num_domains, (int) domains.size());
             for (int i = 0; i < n_domains; i++) {
-                GGML_LOG_INFO("ggml-hex: FASTRPC_GET_DOMAINS[%d]: type %d id %d name '%s' status %d instance-id %d\n",
-                              i, (int) domains[i].type, domains[i].id, domains[i].name, domains[i].status, domains[i].instance_id);
+                HEX_VERBOSE("ggml-hex: FASTRPC_GET_DOMAINS[%d]: type %d id %d name '%s' status %d instance-id %d\n",
+                            i, (int) domains[i].type, domains[i].id, domains[i].name, domains[i].status, domains[i].instance_id);
                 if (domains[i].type != FASTRPC_NSP) {
                     GGML_LOG_DEBUG("ggml-hex:   skipping non-CDSP domain (type=%d)\n", (int) domains[i].type);
                     continue;
@@ -3368,10 +3410,10 @@ static void ggml_hexagon_discover_devices() {
                               domains[i].instance_id, domains[i].id, domains[i].name);
             }
         } else {
-            GGML_LOG_WARN("ggml-hex: FASTRPC_GET_DOMAINS fetch failed (0x%x), using static CDSP domains\n", (unsigned) err);
+            HEX_VERBOSE("ggml-hex: FASTRPC_GET_DOMAINS fetch failed (0x%x), using static CDSP domains\n", (unsigned) err);
         }
     } else if (err != AEE_SUCCESS) {
-        GGML_LOG_DEBUG("ggml-hex: FASTRPC_GET_DOMAINS query failed (0x%x), using static CDSP domains\n", (unsigned) err);
+        HEX_VERBOSE("ggml-hex: FASTRPC_GET_DOMAINS query failed (0x%x), using static CDSP domains\n", (unsigned) err);
     }
 
     // Populate domain IDs and names for all configured devices

ggml/src/ggml-hexagon/htp-common.h

@@ -97,7 +97,7 @@ struct ggml_hexagon_session;
 
 const char * htp_session_name(const struct ggml_hexagon_session * sess);
 size_t       htp_session_vtcm_size(const struct ggml_hexagon_session * sess);
-uint32_t     htp_session_n_threads(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_n_hvx(const struct ggml_hexagon_session * sess);
 uint32_t     htp_session_n_hmx(const struct ggml_hexagon_session * sess);
 uint32_t     htp_session_mdev_count(const struct ggml_hexagon_session * sess);
 

ggml/src/ggml-hexagon/htp-drv.cpp

@@ -1,7 +1,9 @@
+#include <algorithm>
 #include <filesystem>
 #include <set>
 #include <sstream>
 #include <string>
+#include <vector>
 
 #ifdef _WIN32
 #   define WIN32_LEAN_AND_MEAN
@@ -28,6 +30,13 @@
 
 #include <domain.h>
 
+extern int opt_verbose;
+
+#ifndef HEX_VERBOSE
+#define HEX_VERBOSE(...) \
+    if (opt_verbose) GGML_LOG_DEBUG(__VA_ARGS__)
+#endif
+
 //
 // Driver API types
 //
@@ -423,3 +432,188 @@ int htpdrv_get_arch(int domain, int * arch) {
     *arch = (int) ((val >> 4) * 10 + (val & 0x0f));
     return 0;
 }
+
+// Directory of the HTP library as a path with a trailing slash, or empty if it cannot be resolved.
+// Checks GGML_HEXAGON_LIBHTP_DIR first, then the module directory and ../lib.
+// On Windows, non-ASCII directories or paths with spaces use the 8.3 short name.
+std::string htp_lib_dir() {
+    namespace fs = std::filesystem;
+
+    std::vector<fs::path> candidates;
+
+    const char * env_dir = getenv("GGML_HEXAGON_LIBHTP_DIR");
+    if (env_dir && *env_dir) {
+        std::error_code ec;
+        fs::path p = fs::absolute(env_dir, ec);
+        if (!ec) {
+            candidates.push_back(p);
+        }
+    }
+
+#ifdef _WIN32
+    HMODULE hm = nullptr;
+    if (GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+                           (LPCWSTR) &htp_lib_dir, &hm)) {
+        std::wstring buf(32768, L'\0');
+        DWORD n = GetModuleFileNameW(hm, buf.data(), (DWORD) buf.size());
+        if (n > 0 && n < buf.size()) {
+            buf.resize(n);
+            std::error_code ec;
+            fs::path mod_path = fs::canonical(buf, ec);
+            if (ec) {
+                mod_path = fs::absolute(buf, ec);
+            }
+            const fs::path mod_dir = mod_path.parent_path();
+            candidates.push_back(mod_dir);
+            candidates.push_back(mod_dir / "lib");
+            candidates.push_back(mod_dir / ".." / "lib");
+        }
+    }
+#elif defined(__linux__) || defined(__ANDROID__)
+    Dl_info info;
+    if (dladdr((void *) &htp_lib_dir, &info) && info.dli_fname) {
+        std::error_code ec;
+        fs::path mod_path = fs::canonical(info.dli_fname, ec);
+        if (ec) {
+            mod_path = fs::absolute(info.dli_fname, ec);
+        }
+        const fs::path mod_dir = mod_path.parent_path();
+        candidates.push_back(mod_dir);
+        candidates.push_back(mod_dir / "lib");
+        candidates.push_back(mod_dir / ".." / "lib");
+    }
+#endif
+
+    for (size_t i = 0; i < candidates.size(); ++i) {
+        const fs::path & d = candidates[i];
+        std::error_code ec;
+        if (!fs::is_directory(d, ec)) {
+            continue;
+        }
+
+        // For auto-discovered directories, verify it contains an HTP library
+        const bool is_env = (env_dir && i == 0);
+        if (!is_env) {
+            bool found = false;
+            for (const auto & entry : fs::directory_iterator(d, ec)) {
+                if (ec) {
+                    break;
+                }
+                const std::string fn = entry.path().filename().string();
+                if (fn.rfind("libggml-htp-", 0) == 0) {
+                    found = true;
+                    break;
+                }
+            }
+            if (!found) {
+                continue;
+            }
+        }
+
+        std::error_code ec_canon;
+        fs::path abs_d = fs::canonical(d, ec_canon);
+        if (ec_canon) {
+            abs_d = fs::absolute(d, ec_canon);
+        }
+
+#ifdef _WIN32
+        std::wstring dir = abs_d.wstring();
+        auto is_ascii = [](const std::wstring & s) {
+            return std::all_of(s.begin(), s.end(), [](wchar_t c) { return c < 0x80; });
+        };
+        auto has_spaces = [](const std::wstring & s) {
+            return std::any_of(s.begin(), s.end(), [](wchar_t c) { return c == L' '; });
+        };
+
+        if (!is_ascii(dir) || has_spaces(dir)) {
+            DWORD m = GetShortPathNameW(dir.c_str(), nullptr, 0);
+            if (m > 0) {
+                std::wstring sp(m, L'\0');
+                DWORD res = GetShortPathNameW(dir.c_str(), sp.data(), (DWORD) sp.size());
+                if (res > 0 && res < sp.size()) {
+                    sp.resize(res);
+                    if (is_ascii(sp)) {
+                        dir = sp;
+                    }
+                }
+            }
+            if (!is_ascii(dir)) {
+                continue;
+            }
+        }
+        if (dir.back() != L'\\' && dir.back() != L'/') {
+            dir += L'\\';
+        }
+        return std::string(dir.begin(), dir.end());
+#else
+        std::string dir = abs_d.string();
+        if (dir.back() != '/') {
+            dir += '/';
+        }
+        return dir;
+#endif
+    }
+
+    return {};
+}
+
+int htpdrv_set_dsp_search_path(int domain_id, const std::string & path, const char * dev_name) {
+    if (path.empty()) {
+        return AEE_SUCCESS;
+    }
+
+    std::string clean_path = path;
+    while (clean_path.size() > 1 && (clean_path.back() == '/' || clean_path.back() == '\\')) {
+        clean_path.pop_back();
+    }
+
+    fastrpc_manage_dsp_lib_path p {};
+    p.effec_domain_id = domain_id;
+    p.type            = FASTRPC_SET_PATH;
+    p.path            = const_cast<char *>(clean_path.c_str());
+    p.pathlen         = clean_path.size() + 1;
+
+    int err = remote_session_control(FASTRPC_MANAGE_DSP_LIBRARY_PATH, (void *) &p, sizeof(p));
+    if (err != AEE_SUCCESS && domain_id != -1) {
+        p.effec_domain_id = -1;
+        err = remote_session_control(FASTRPC_MANAGE_DSP_LIBRARY_PATH, (void *) &p, sizeof(p));
+    }
+
+    const char * name = dev_name ? dev_name : "driver";
+    if (err == AEE_SUCCESS) {
+        GGML_LOG_INFO("ggml-hex: %s set DSP search path (domain %d) to %s\n", name, p.effec_domain_id, clean_path.c_str());
+        return AEE_SUCCESS;
+    }
+
+    HEX_VERBOSE("ggml-hex: %s FASTRPC_MANAGE_DSP_LIBRARY_PATH failed (0x%x)\n", name, err);
+    return err;
+}
+
+void htpdrv_set_dsp_env_path(const std::string & path, const char * dev_name) {
+    if (path.empty()) {
+        return;
+    }
+
+    std::string clean_path = path;
+    while (clean_path.size() > 1 && (clean_path.back() == '/' || clean_path.back() == '\\')) {
+        clean_path.pop_back();
+    }
+
+    const char * name = dev_name ? dev_name : "driver";
+    HEX_VERBOSE("ggml-hex: %s setting DSP search path via process env to %s\n", name, clean_path.c_str());
+
+    static const char * const env_names[] = {
+        "ADSP_LIBRARY_PATH",
+        "CDSP_LIBRARY_PATH",
+        "CDSP1_LIBRARY_PATH",
+        "CDSP2_LIBRARY_PATH",
+        "CDSP3_LIBRARY_PATH",
+    };
+    for (const char * env_name : env_names) {
+#ifdef _WIN32
+        SetEnvironmentVariableA(env_name, clean_path.c_str());
+#else
+        setenv(env_name, clean_path.c_str(), 1);
+#endif
+    }
+}

ggml/src/ggml-hexagon/htp-drv.h

@@ -120,4 +120,10 @@ HTPDRV_API int remote_system_request(system_req_payload * req);
 
 #ifdef __cplusplus
 }
+
+#include <string>
+
+std::string htp_lib_dir();
+int         htpdrv_set_dsp_search_path(int domain_id, const std::string & path, const char * dev_name = nullptr);
+void        htpdrv_set_dsp_env_path(const std::string & path, const char * dev_name = nullptr);
 #endif

ggml/src/ggml-hexagon/htp-fusion.cpp

@@ -239,11 +239,11 @@ static bool fuse_mul_mat_add(const ggml_hexagon_session * sess, htp_opnode & pre
     if (kparams.n_hmx > 0 && orig_kparams->n_hmx > 0) {
         if (kparams.m_chunk < orig_kparams->m_chunk ||
             kparams.n_chunk < orig_kparams->n_chunk ||
-            kparams.n_act_threads < orig_kparams->n_act_threads) {
-            HEX_VERBOSE("ggml-hex: %s skip MUL_MAT_ADD fusion: HMX efficiency reduced (m %d->%d, n %d->%d, th %d->%d)\n",
+            kparams.n_act_hvx < orig_kparams->n_act_hvx) {
+            HEX_VERBOSE("ggml-hex: %s skip MUL_MAT_ADD fusion: HMX efficiency reduced (m %d->%d, n %d->%d, hvx %d->%d)\n",
                         htp_session_name(sess), orig_kparams->m_chunk, kparams.m_chunk,
                         orig_kparams->n_chunk, kparams.n_chunk,
-                        orig_kparams->n_act_threads, kparams.n_act_threads);
+                        orig_kparams->n_act_hvx, kparams.n_act_hvx);
             return false;
         }
     }
@@ -546,8 +546,8 @@ static bool fuse_binary_nx(const ggml_hexagon_session * sess, htp_opnode & prev,
 
         const uint32_t total_elems = (uint32_t) ggml_nelements(src0);
         const uint32_t elem_size   = (uint32_t) ggml_type_size(src0->type);
-        const uint32_t max_threads = htp_session_n_threads(sess);
-        if (!htp_binary_solve_chunked(&L, &kparams, total_elems, elem_size, max_threads, n_srcs, false, vtcm_size)) {
+        const uint32_t max_hvx     = htp_session_n_hvx(sess);
+        if (!htp_binary_solve_chunked(&L, &kparams, total_elems, elem_size, max_hvx, n_srcs, false, vtcm_size)) {
             return false;
         }
     }

ggml/src/ggml-hexagon/htp-opnode.cpp

@@ -157,7 +157,7 @@ static std::string format_kernel_params(const htp_opnode & node) {
         } else if (type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
             path = "hvx-tiled";
         }
-        snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", path, (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
         const auto * kparams = node.kparams<htp_fa_kernel_params>();
         const char * path = "unknown";
@@ -167,10 +167,11 @@ static std::string format_kernel_params(const htp_opnode & node) {
         } else if (type == HTP_FA_KERNEL_HVX) {
             path = "hvx";
         }
-        snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", path, (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (htp_op_is_unary(node.opcode)) {
         const auto * kparams = node.kparams<htp_unary_kernel_params>();
-        snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", kparams->col_tile ? "wide-row" : "row-block",
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_MDEV_GROUP && node.node) {
         snprintf(str, max_size, "idx %d count %d", (int) node.node->op_params[0], (int) node.dst()->ne[1]);
     } else if ((node.opcode == HTP_OP_FENCE || node.opcode == HTP_OP_CPY_FENCE) && node.node) {
@@ -179,41 +180,43 @@ static std::string format_kernel_params(const htp_opnode & node) {
         snprintf(str, max_size, "seq 0x%x -> 0x%x", (uint32_t) node.node->op_params[0], (uint32_t) node.node->op_params[1]);
     } else if (node.opcode == HTP_OP_SSM_CONV) {
         const auto * kparams = node.kparams<htp_ssm_conv_kernel_params>();
-        snprintf(str, max_size, "%s vtcm %d", kparams->n_t == 1 ? "decode" : "prefill", (int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", kparams->n_t == 1 ? "decode" : "prefill",
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_SOFTMAX) {
         const auto * kparams = node.kparams<htp_softmax_kernel_params>();
-        snprintf(str, max_size, "k%d nth %d vtcm %d", (int) kparams->kernel_type, (int) kparams->n_threads, (int) kparams->vtcm_size);
+        snprintf(str, max_size, "k%d nhvx %u vtcm %u", (int) kparams->kernel_type,
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_GATED_DELTA_NET) {
         const auto * kparams = node.kparams<htp_gdn_kernel_params>();
         const char * path = (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) ? "hmx-chunked" : "hvx-recurrent";
-        snprintf(str, max_size, "%s-%s vtcm %u", path, kparams->kda ? "kda" : "scalar",
-                 (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s-%s nhvx %u vtcm %u", path, kparams->kda ? "kda" : "scalar",
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_DSV4_HC_PRE) {
         const auto * kparams = node.kparams<htp_dsv4_hc_pre_kernel_params>();
         const char * ktype_str = kparams->kernel_type == HTP_DSV4_HC_PRE_KERNEL_GATED ? "gated" : "plain";
-        snprintf(str, max_size, "%s nth %u vtcm %u", ktype_str,
-                 (unsigned int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", ktype_str,
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_DSV4_HC_POST) {
         const auto * kparams = node.kparams<htp_dsv4_hc_post_kernel_params>();
         const char * ktype_str = kparams->kernel_type == HTP_DSV4_HC_POST_KERNEL_COMB ? "comb" : "identity";
-        snprintf(str, max_size, "%s nth %u vtcm %u", ktype_str,
-                 (unsigned int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", ktype_str,
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_LIGHTNING_INDEXER) {
         const auto * kparams = node.kparams<htp_lid_kernel_params>();
-        snprintf(str, max_size, "%s tok %u bc %u nth %u vtcm %u",
+        snprintf(str, max_size, "%s nhvx %u tok %u bc %u vtcm %u",
                  kparams->kernel_type == HTP_LID_KERNEL_HMX ? "hmx" : "hvx",
-                 (unsigned int) kparams->n_tokens_tile, (unsigned int) kparams->kv_block,
-                 (unsigned int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->n_tokens_tile,
+                 (unsigned int) kparams->kv_block, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_MUL || node.opcode == HTP_OP_ADD || node.opcode == HTP_OP_ADD_ID ||
                node.opcode == HTP_OP_SUB || node.opcode == HTP_OP_DIV) {
         const auto * kparams = node.kparams<htp_binary_kernel_params>();
-        snprintf(str, max_size, "vtcm %u", (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "nhvx %u vtcm %u", (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
         const auto * kparams = node.kparams<htp_sort_kernel_params>();
-        snprintf(str, max_size, "%s nth %d nchk %d chk %d vtcm %d",
+        snprintf(str, max_size, "%s nhvx %u nchk %u chk %u vtcm %u",
                  node.opcode == HTP_OP_TOP_K ? "top_k" : "argsort",
-                 (int) kparams->n_threads, (int) kparams->n_chunks,
-                 (int) kparams->chunk_elems, (int) kparams->vtcm_size);
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->n_chunks,
+                 (unsigned int) kparams->chunk_elems, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_GET_ROWS) {
         const auto * kparams = node.kparams<htp_get_rows_kernel_params>();
         const char * ktype_str = "unknown";
@@ -222,16 +225,16 @@ static std::string format_kernel_params(const htp_opnode & node) {
             case HTP_GET_ROWS_KERNEL_TILED:    ktype_str = "tiled"; break;
             case HTP_GET_ROWS_KERNEL_FLAT:     ktype_str = "flat"; break;
         }
-        snprintf(str, max_size, "%s%s vtcm %u", ktype_str, kparams->n_threads > 1 ? "-multi" : "", (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", ktype_str, (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (htp_op_is_glu(node.opcode)) {
         const auto * kparams = node.kparams<htp_glu_kernel_params>();
-        snprintf(str, max_size, "nth %d vtcm %u", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "nhvx %u vtcm %u", (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_PAD) {
         const auto * kparams = node.kparams<htp_pad_kernel_params>();
-        snprintf(str, max_size, "k%d nth %d vtcm %u", (int) kparams->kernel_type, (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "k%d nhvx %u vtcm %u", (int) kparams->kernel_type, (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_CUMSUM) {
         const auto * kparams = node.kparams<htp_cumsum_kernel_params>();
-        snprintf(str, max_size, "%s nth %d vtcm %u", kparams->use_dma ? "dma" : "hvx", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", kparams->use_dma ? "dma" : "hvx", (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_TOP_K_MASK) {
         const auto * kparams = node.kparams<htp_top_k_mask_kernel_params>();
         const char * ktype_str = "?";
@@ -239,19 +242,19 @@ static std::string format_kernel_params(const htp_opnode & node) {
             case HTP_TOP_K_MASK_KERNEL_ROWS:  ktype_str = "rows"; break;
             case HTP_TOP_K_MASK_KERNEL_PICKS: ktype_str = "picks"; break;
         }
-        snprintf(str, max_size, "%s nth %d tile %u x%u vtcm %u", ktype_str, (int) kparams->n_threads, (unsigned int) kparams->tile_cols, (unsigned int) kparams->n_col_tiles, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u tile %u x%u vtcm %u", ktype_str, (unsigned int) kparams->n_hvx, (unsigned int) kparams->tile_cols, (unsigned int) kparams->n_col_tiles, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_FILL) {
         const auto * kparams = node.kparams<htp_fill_kernel_params>();
-        snprintf(str, max_size, "%s nth %d", kparams->kernel_type == HTP_FILL_KERNEL_CONTIGUOUS ? "contig" : "strided", (int) kparams->n_threads);
+        snprintf(str, max_size, "%s nhvx %u", kparams->kernel_type == HTP_FILL_KERNEL_CONTIGUOUS ? "contig" : "strided", (unsigned int) kparams->n_hvx);
     } else if (node.opcode == HTP_OP_ROLL) {
         const auto * kparams = node.kparams<htp_roll_kernel_params>();
-        snprintf(str, max_size, "k%d nth %d", (int) kparams->kernel_type, (int) kparams->n_threads);
+        snprintf(str, max_size, "k%d nhvx %u", (int) kparams->kernel_type, (unsigned int) kparams->n_hvx);
     } else if (node.opcode == HTP_OP_DIAG) {
         const auto * kparams = node.kparams<htp_diag_kernel_params>();
-        snprintf(str, max_size, "%s nth %d vtcm %u", kparams->use_dma ? "dma" : "hvx", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", kparams->use_dma ? "dma" : "hvx", (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_SOLVE_TRI) {
         const auto * kparams = node.kparams<htp_solve_tri_kernel_params>();
-        snprintf(str, max_size, "%s nth %d", kparams->batched ? "batched" : "chunked", (int) kparams->n_threads);
+        snprintf(str, max_size, "%s nhvx %u", kparams->batched ? "batched" : "chunked", (unsigned int) kparams->n_hvx);
     } else if (node.opcode == HTP_OP_UPSCALE) {
         const auto * kparams = node.kparams<htp_upscale_kernel_params>();
         const char * ktype_str = "unknown";
@@ -264,7 +267,7 @@ static std::string format_kernel_params(const htp_opnode & node) {
             case HTP_UPSCALE_KERNEL_BILINEAR_AA: ktype_str = "bilinear-aa"; break;
             case HTP_UPSCALE_KERNEL_BICUBIC:     ktype_str = "bicubic";     break;
         }
-        snprintf(str, max_size, "%s nth %d chk %d vtcm %u", ktype_str, (int) kparams->n_threads, (int) kparams->dst_nrows_per_chunk, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u chk %u vtcm %u", ktype_str, (unsigned int) kparams->n_hvx, (unsigned int) kparams->dst_nrows_per_chunk, (unsigned int) kparams->vtcm_size);
     } else {
         snprintf(str, max_size, "----");
     }

ggml/src/ggml-hexagon/htp/allreduce-ops.c

@@ -15,7 +15,7 @@
 #include "hvx-utils.h"
 #include "htp-tensor.h"
 #include "dma-queue.h"
-#include "hex-profile.h"
+#include "htp-trace.h"
 #include "allreduce-ops.h"
 #include "htp-fence.h"
 
@@ -29,9 +29,9 @@ struct htp_allreduce_context {
     uint32_t row_size_aligned;
     uint32_t rank_elem_start;
     uint32_t rank_nelem;
-    uint32_t elems_per_thread;
+    uint32_t elems_per_hvx;
     uint32_t block_elems;
-    uint32_t vtcm_size_per_thread;
+    uint32_t vtcm_size_per_hvx;
     bool     is_row_bcast;
     uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];
     uint8_t * dst_spad_base;
@@ -39,7 +39,8 @@ struct htp_allreduce_context {
 };
 
 #define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD)                                         \
-static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                   \
+static void allreduce_thread_dma_1d_##SUFFIX(const htp_thread_t * th, unsigned int nth, unsigned int ith,         \
+                                             void * data) {                                                       \
     struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data;                                  \
     struct htp_ops_context * octx = actx->octx;                                                                   \
                                                                                                                   \
@@ -47,19 +48,19 @@ static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith,
     const uint32_t n_dsts      = actx->n_dsts;                                                                    \
     const uint32_t block_elems = actx->block_elems;                                                               \
                                                                                                                   \
-    const uint32_t dr  = actx->elems_per_thread;                                                                  \
+    const uint32_t dr  = actx->elems_per_hvx;                                                                     \
     const uint32_t ir0 = actx->rank_elem_start + dr * ith;                                                        \
     const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem);                                 \
     if (ir0 >= ir1) return;                                                                                       \
                                                                                                                   \
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                        \
-    dma_queue * dma_q = octx->ctx->dma[ith];                                                                      \
+    struct htp_thread_trace * tr = th->trace;                                                                     \
+    dma_queue * dma_q = th->dma_q;                                                                                \
                                                                                                                   \
-    const size_t vtcm_thread_offset = ith * actx->vtcm_size_per_thread;                                           \
-    uint8_t * dst_spad_base = actx->dst_spad_base + vtcm_thread_offset;                                           \
-    uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + vtcm_thread_offset) : NULL;                        \
+    const size_t vtcm_hvx_offset = ith * actx->vtcm_size_per_hvx;                                                 \
+    uint8_t * dst_spad_base = actx->dst_spad_base + vtcm_hvx_offset;                                              \
+    uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + vtcm_hvx_offset) : NULL;                           \
                                                                                                                   \
-    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                      \
+    const size_t spad_half = actx->vtcm_size_per_hvx / 2;                                                         \
     uint32_t ir_prefetch = ir0;                                                                                   \
     int spad_idx = 0;                                                                                             \
                                                                                                                   \
@@ -72,7 +73,7 @@ static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith,
             dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0);              \
         }                                                                                                         \
         for (uint32_t s = 0; s < n_ranks; s++) {                                                                  \
-            uint8_t * s_spad = actx->src_spad_base[s] + vtcm_thread_offset + spad_idx * spad_half;                \
+            uint8_t * s_spad = actx->src_spad_base[s] + vtcm_hvx_offset + spad_idx * spad_half;                   \
             const dma_addr_t s_ddr = octx->src[s]->data + ir_prefetch * sizeof(TYPE);                             \
             dma_queue_push(dma_q, dma_make_data(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1);              \
         }                                                                                                         \
@@ -99,13 +100,13 @@ static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith,
             dma_queue_pop(dma_q);                                                                                 \
         }                                                                                                         \
         uint8_t * d_spad  = dst_spad_base + comp_spad_idx * spad_half;                                            \
-        uint8_t * s0_spad = actx->src_spad_base[0] + vtcm_thread_offset + comp_spad_idx * spad_half;              \
-        uint8_t * s1_spad = actx->src_spad_base[1] + vtcm_thread_offset + comp_spad_idx * spad_half;              \
+        uint8_t * s0_spad = actx->src_spad_base[0] + vtcm_hvx_offset + comp_spad_idx * spad_half;                 \
+        uint8_t * s1_spad = actx->src_spad_base[1] + vtcm_hvx_offset + comp_spad_idx * spad_half;                 \
         uint8_t * r_spad  = HAS_ADD ? (res_spad_base + comp_spad_idx * spad_half) : NULL;                         \
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                         \
         HVX_ADD_FN(d_spad, s0_spad, s1_spad, cur_elems);                                                          \
         for (uint32_t s = 2; s < n_ranks; s++) {                                                                  \
-            uint8_t * ss_spad = actx->src_spad_base[s] + vtcm_thread_offset + comp_spad_idx * spad_half;          \
+            uint8_t * ss_spad = actx->src_spad_base[s] + vtcm_hvx_offset + comp_spad_idx * spad_half;             \
             HVX_ADD_FN(d_spad, d_spad, ss_spad, cur_elems);                                                       \
         }                                                                                                         \
         if (HAS_ADD) {                                                                                            \
@@ -120,7 +121,7 @@ static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith,
             uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch);                                            \
             size_t   next_bytes = next_elems * sizeof(TYPE);                                                      \
             for (uint32_t s = 0; s < n_ranks; s++) {                                                              \
-                uint8_t * s_spad = actx->src_spad_base[s] + vtcm_thread_offset + comp_spad_idx * spad_half;       \
+                uint8_t * s_spad = actx->src_spad_base[s] + vtcm_hvx_offset + comp_spad_idx * spad_half;          \
                 const dma_addr_t s_next = octx->src[s]->data + ir_prefetch * sizeof(TYPE);                        \
                 dma_queue_push(dma_q, dma_make_data(s_spad, s_next), next_bytes, next_bytes, next_bytes, 1);      \
             }                                                                                                     \
@@ -143,7 +144,7 @@ DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f16, __fp16, hvx_add_f16_aaa, 1)
 DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f32, float,  hvx_add_f32_aaa, 1)
 
 #define DEFINE_ALLREDUCE_THREAD_DMA_2D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD, IS_ROW_BCAST)                                                           \
-static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                                                   \
+static void allreduce_thread_dma_2d_##SUFFIX(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {                          \
     struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data;                                                                  \
     struct htp_ops_context * octx = actx->octx;                                                                                                   \
                                                                                                                                                   \
@@ -154,19 +155,19 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
     const uint32_t row_size_aligned = actx->row_size_aligned;                                                                                     \
     const uint32_t row_bytes        = ne0 * sizeof(TYPE);                                                                                         \
                                                                                                                                                   \
-    const uint32_t dr  = actx->elems_per_thread;                                                                                                  \
+    const uint32_t dr  = actx->elems_per_hvx;                                                                                                     \
     const uint32_t r0  = actx->rank_elem_start + dr * ith;                                                                                        \
     const uint32_t r1  = MIN(r0 + dr, actx->rank_elem_start + actx->rank_nelem);                                                                  \
     if (r0 >= r1) return;                                                                                                                         \
                                                                                                                                                   \
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                                        \
-    dma_queue * dma_q = octx->ctx->dma[ith];                                                                                                      \
+    struct htp_thread_trace * tr = th->trace;                                                                                                     \
+    dma_queue * dma_q = th->dma_q;                                                                                                                \
                                                                                                                                                   \
-    const size_t vtcm_thread_offset = ith * actx->vtcm_size_per_thread;                                                                           \
-    uint8_t * dst_spad_base = actx->dst_spad_base + vtcm_thread_offset;                                                                           \
-    uint8_t * res_spad_base = HAS_ADD ? (IS_ROW_BCAST ? actx->res_spad_base : (actx->res_spad_base + vtcm_thread_offset)) : NULL;                 \
+    const size_t vtcm_hvx_offset = ith * actx->vtcm_size_per_hvx;                                                                                 \
+    uint8_t * dst_spad_base = actx->dst_spad_base + vtcm_hvx_offset;                                                                              \
+    uint8_t * res_spad_base = HAS_ADD ? (IS_ROW_BCAST ? actx->res_spad_base : (actx->res_spad_base + vtcm_hvx_offset)) : NULL;                    \
                                                                                                                                                   \
-    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                                                      \
+    const size_t spad_half = actx->vtcm_size_per_hvx / 2;                                                                                         \
     uint32_t r_prefetch = r0;                                                                                                                     \
     int spad_idx = 0;                                                                                                                             \
                                                                                                                                                   \
@@ -178,7 +179,7 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
             dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, 0);                            \
         }                                                                                                                                         \
         for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                  \
-            uint8_t * s_spad = actx->src_spad_base[s] + vtcm_thread_offset + spad_idx * spad_half;                                                \
+            uint8_t * s_spad = actx->src_spad_base[s] + vtcm_hvx_offset + spad_idx * spad_half;                                                   \
             const dma_addr_t s_ddr = octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                                       \
             dma_queue_push(dma_q, dma_make_data(s_spad, s_ddr), row_size_aligned, octx->src[s]->nb[1], row_bytes, cur_rows);                      \
         }                                                                                                                                         \
@@ -204,8 +205,8 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
             dma_queue_pop(dma_q);                                                                                                                 \
         }                                                                                                                                         \
         uint8_t * d_spad  = dst_spad_base + comp_spad_idx * spad_half;                                                                            \
-        uint8_t * s0_spad = actx->src_spad_base[0] + vtcm_thread_offset + comp_spad_idx * spad_half;                                              \
-        uint8_t * s1_spad = actx->src_spad_base[1] + vtcm_thread_offset + comp_spad_idx * spad_half;                                              \
+        uint8_t * s0_spad = actx->src_spad_base[0] + vtcm_hvx_offset + comp_spad_idx * spad_half;                                                 \
+        uint8_t * s1_spad = actx->src_spad_base[1] + vtcm_hvx_offset + comp_spad_idx * spad_half;                                                 \
         uint8_t * r_spad  = (HAS_ADD && !IS_ROW_BCAST) ? (res_spad_base + comp_spad_idx * spad_half) : NULL;                                      \
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);                                                                          \
         for (uint32_t row = 0; row < cur_rows; row++) {                                                                                           \
@@ -214,7 +215,7 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
             const uint8_t * s1_row = s1_spad + row * row_size_aligned;                                                                            \
             HVX_ADD_FN(d_row, s0_row, s1_row, ne0);                                                                                               \
             for (uint32_t s = 2; s < n_ranks; s++) {                                                                                              \
-                const uint8_t * ss_row = actx->src_spad_base[s] + vtcm_thread_offset + comp_spad_idx * spad_half + row * row_size_aligned;        \
+                const uint8_t * ss_row = actx->src_spad_base[s] + vtcm_hvx_offset + comp_spad_idx * spad_half + row * row_size_aligned;           \
                 HVX_ADD_FN(d_row, d_row, ss_row, ne0);                                                                                            \
             }                                                                                                                                     \
             if (HAS_ADD) {                                                                                                                        \
@@ -230,7 +231,7 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
         if (r_prefetch < r1) {                                                                                                                    \
             uint32_t next_rows = MIN(block_rows, r1 - r_prefetch);                                                                                \
             for (uint32_t s = 0; s < n_ranks; s++) {                                                                                              \
-                uint8_t * s_spad = actx->src_spad_base[s] + vtcm_thread_offset + comp_spad_idx * spad_half;                                       \
+                uint8_t * s_spad = actx->src_spad_base[s] + vtcm_hvx_offset + comp_spad_idx * spad_half;                                          \
                 const dma_addr_t s_next = octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                                  \
                 dma_queue_push(dma_q, dma_make_data(s_spad, s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows);                \
             }                                                                                                                                     \
@@ -259,19 +260,19 @@ static int validate_allreduce(
     const struct htp_allreduce_kernel_params * kparams,
     uint32_t n_ranks
 ) {
-    if (kparams->n_threads == 0 || (uint32_t) kparams->n_threads > octx->ctx->n_threads) {
+    if (kparams->n_hvx == 0 || (uint32_t) kparams->n_hvx > octx->ctx->n_hvx) {
         return HTP_STATUS_INVAL_PARAMS;
     }
-    octx->n_threads = (uint32_t) kparams->n_threads;
+    octx->n_hvx = (uint32_t) kparams->n_hvx;
 
-    if (kparams->vtcm_size_per_thread <= 0 || kparams->vtcm_size <= 0) {
+    if (kparams->vtcm_size_per_hvx <= 0 || kparams->vtcm_size <= 0) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
     const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
     const size_t n_vtcm_buffers = htp_allreduce_vtcm_buffer_count(
-        n_ranks, octx->n_threads, has_add, kparams->is_row_bcast != 0);
-    const size_t vtcm_size = n_vtcm_buffers * (size_t) kparams->vtcm_size_per_thread;
+        n_ranks, octx->n_hvx, has_add, kparams->is_row_bcast != 0);
+    const size_t vtcm_size = n_vtcm_buffers * (size_t) kparams->vtcm_size_per_hvx;
     if (vtcm_size != (size_t) kparams->vtcm_size) {
         return HTP_STATUS_INVAL_PARAMS;
     }
@@ -360,10 +361,10 @@ int op_allreduce(struct htp_ops_context * octx) {
 
     // 2. Multi-threaded Reduction across assigned rank chunk
     if (nelem > 0) {
-        const uint32_t n_threads            = (uint32_t) kparams->n_threads;
+        const uint32_t n_hvx                = (uint32_t) kparams->n_hvx;
         const uint32_t block_elems          = (uint32_t) kparams->block_elems;
-        const uint32_t elems_per_thread     = (uint32_t) kparams->elems_per_thread;
-        const uint32_t vtcm_size_per_thread = (uint32_t) kparams->vtcm_size_per_thread;
+        const uint32_t elems_per_hvx        = (uint32_t) kparams->elems_per_hvx;
+        const uint32_t vtcm_size_per_hvx    = (uint32_t) kparams->vtcm_size_per_hvx;
 
         struct htp_allreduce_context actx;
         actx.octx                 = octx;
@@ -375,9 +376,9 @@ int op_allreduce(struct htp_ops_context * octx) {
         actx.row_size_aligned     = (uint32_t) kparams->row_size_aligned;
         actx.rank_elem_start      = (uint32_t) kparams->rank_elem_start;
         actx.rank_nelem           = (uint32_t) kparams->rank_nelem;
-        actx.elems_per_thread     = elems_per_thread;
+        actx.elems_per_hvx        = elems_per_hvx;
         actx.block_elems          = block_elems;
-        actx.vtcm_size_per_thread = vtcm_size_per_thread;
+        actx.vtcm_size_per_hvx    = vtcm_size_per_hvx;
         actx.is_row_bcast         = (kparams->is_row_bcast != 0);
 
         work_queue_func_t reduce_fun = NULL;
@@ -409,13 +410,13 @@ int op_allreduce(struct htp_ops_context * octx) {
         uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
         for (uint32_t s = 0; s < n_ranks; s++) {
             actx.src_spad_base[s] = vtcm_ptr;
-            vtcm_ptr += n_threads * vtcm_size_per_thread;
+            vtcm_ptr += n_hvx * vtcm_size_per_hvx;
         }
         actx.dst_spad_base = vtcm_ptr;
-        vtcm_ptr += n_threads * vtcm_size_per_thread;
+        vtcm_ptr += n_hvx * vtcm_size_per_hvx;
         if (has_add) {
             actx.res_spad_base = vtcm_ptr;
-            vtcm_ptr += (actx.is_row_bcast ? 1 : n_threads) * vtcm_size_per_thread;
+            vtcm_ptr += (actx.is_row_bcast ? 1 : n_hvx) * vtcm_size_per_hvx;
         }
 
         if (has_add && actx.is_row_bcast) {
@@ -426,7 +427,7 @@ int op_allreduce(struct htp_ops_context * octx) {
             dma_queue_pop(dma_q);
         }
 
-        work_queue_run(octx->ctx->work_queue, reduce_fun, &actx, n_threads);
+        work_queue_run(octx->ctx->hvx_queue, reduce_fun, &actx, n_hvx);
     }
 
     // 4. Exit Barrier: Synchronize all ranks after writing

ggml/src/ggml-hexagon/htp/allreduce-ops.h

@@ -24,21 +24,21 @@ enum htp_allreduce_mode {
 
 static inline size_t htp_allreduce_vtcm_buffer_count(
     uint32_t n_ranks,
-    uint32_t n_threads,
+    uint32_t n_hvx,
     bool has_add,
     bool is_row_bcast
 ) {
-    return (size_t) (n_ranks + 1) * n_threads + (has_add ? (is_row_bcast ? 1 : n_threads) : 0);
+    return (size_t) (n_ranks + 1) * n_hvx + (has_add ? (is_row_bcast ? 1 : n_hvx) : 0);
 }
 
 struct htp_allreduce_kernel_params {
     uint8_t kernel_type;
     int32_t rank;
     int32_t n_ranks;
-    int32_t n_threads;
+    int32_t n_hvx;
     int32_t block_elems;          // 1D: block_elems, 2D: block_rows
-    int32_t elems_per_thread;     // 1D: nelem_per_thread, 2D: nrows_per_thread
-    int32_t vtcm_size_per_thread;
+    int32_t elems_per_hvx;        // 1D: nelem_per_hvx, 2D: nrows_per_hvx
+    int32_t vtcm_size_per_hvx;
     int32_t vtcm_size;
     int32_t ne0;
     int32_t ne1;

ggml/src/ggml-hexagon/htp/argsort-ops.c

@@ -463,10 +463,10 @@ struct htp_sort_chunk_ctx {
     size_t                                phase1_slot_size;
     float *                               merge_values;
     int32_t *                             merge_indices;
-    uint32_t                              real_count[HTP_MAX_NTHREADS];
+    uint32_t                              real_count[HTP_OP_MAX_NHVX];
 };
 
-static void htp_sort_chunk_job(unsigned int n, unsigned int i, void * data) {
+static void htp_sort_chunk_job(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     struct htp_sort_chunk_ctx * cctx = (struct htp_sort_chunk_ctx *) data;
     struct htp_ops_context * octx = cctx->octx;
     const struct htp_sort_kernel_params * kparams = cctx->kparams;
@@ -487,7 +487,7 @@ static void htp_sort_chunk_job(unsigned int n, unsigned int i, void * data) {
     float *   merge_values  = cctx->merge_values;
     int32_t * merge_indices = cctx->merge_indices;
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     const dma_addr_t src_addr = src0->data + (size_t) chunk_base * sizeof(float);
 
     if (kparams->is_top_k) {
@@ -503,7 +503,7 @@ static void htp_sort_chunk_job(unsigned int n, unsigned int i, void * data) {
         }
 
         if (local_k == 1 && ne00 >= 128*1024) {
-            dma_sync_read(octx->ctx->dma[i], values_buf, src_addr, real_count * sizeof(float));
+            dma_sync_read(th->dma_q, values_buf, src_addr, real_count * sizeof(float));
 
             htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
             float max_value;
@@ -515,7 +515,7 @@ static void htp_sort_chunk_job(unsigned int n, unsigned int i, void * data) {
         }
 
         if (local_k > 1 && local_k <= 64 && chunk_elems > 1024) {
-            dma_sync_read(octx->ctx->dma[i], values_buf, src_addr, real_count * sizeof(float));
+            dma_sync_read(th->dma_q, values_buf, src_addr, real_count * sizeof(float));
 
             htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
             top_k_scan(values_buf, real_count, local_k, &merge_values[i * local_k], &merge_indices[i * local_k]);
@@ -526,7 +526,7 @@ static void htp_sort_chunk_job(unsigned int n, unsigned int i, void * data) {
             return;
         }
 
-        dma_sync_read(octx->ctx->dma[i], values_buf, src_addr, real_count * sizeof(float));
+        dma_sync_read(th->dma_q, values_buf, src_addr, real_count * sizeof(float));
 
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
         if (chunk_elems > real_count) {
@@ -542,7 +542,7 @@ static void htp_sort_chunk_job(unsigned int n, unsigned int i, void * data) {
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
     } else {
         if (real_count > 0) {
-            dma_sync_read(octx->ctx->dma[i], values_buf, src_addr, real_count * sizeof(float));
+            dma_sync_read(th->dma_q, values_buf, src_addr, real_count * sizeof(float));
 
             htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
             init_indices_ramp(indices_buf, real_count, chunk_base);
@@ -562,7 +562,7 @@ struct htp_argsort_merge_l1_ctx {
     bool                        asc;
 };
 
-static void htp_argsort_merge_l1_job(unsigned int n, unsigned int i, void * data) {
+static void htp_argsort_merge_l1_job(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     struct htp_argsort_merge_l1_ctx * mctx = (struct htp_argsort_merge_l1_ctx *) data;
     struct htp_sort_chunk_ctx * cctx = mctx->cctx;
     const struct htp_sort_kernel_params * kparams = cctx->kparams;
@@ -585,7 +585,7 @@ static void htp_argsort_merge_l1_job(unsigned int n, unsigned int i, void * data
     float *   out_val = cctx->merge_values  + out_offset;
     int32_t * out_idx = cctx->merge_indices + out_offset;
 
-    struct htp_thread_trace * tr = &cctx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
     merge_runs(val0, idx0, count0, val1, idx1, count1, out_val, out_idx, mctx->asc);
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
@@ -596,7 +596,7 @@ struct htp_argsort_merge_l2_ctx {
     bool                        asc;
 };
 
-static void htp_argsort_merge_l2_job(unsigned int n, unsigned int i, void * data) {
+static void htp_argsort_merge_l2_job(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     struct htp_argsort_merge_l2_ctx * mctx = (struct htp_argsort_merge_l2_ctx *) data;
     struct htp_sort_chunk_ctx * cctx = mctx->cctx;
     const struct htp_sort_kernel_params * kparams = cctx->kparams;
@@ -619,7 +619,7 @@ static void htp_argsort_merge_l2_job(unsigned int n, unsigned int i, void * data
     float *   out_val   = (float *) slot_base;
     int32_t * out_idx   = (int32_t *) (slot_base + 4 * chunk_elems * sizeof(float));
 
-    struct htp_thread_trace * tr = &cctx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
     merge_runs(val0, idx0, count0, val1, idx1, count1, out_val, out_idx, mctx->asc);
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
@@ -632,7 +632,7 @@ struct htp_sort_multi_row_ctx {
     size_t                                phase1_slot_size;
     uint32_t                              row_start;
     uint32_t                              row_end;
-    uint32_t                              nrows_per_thread;
+    uint32_t                              nrows_per_hvx;
 };
 
 struct htp_tensor_row_coord {
@@ -721,7 +721,7 @@ static inline void compute_row_sort(
     }
 }
 
-static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data) {
+static void htp_sort_multi_row_job(const htp_thread_t * th, unsigned int n, unsigned int ith, void * data) {
     struct htp_sort_multi_row_ctx * mctx = (struct htp_sort_multi_row_ctx *) data;
     struct htp_ops_context * octx = mctx->octx;
     const struct htp_sort_kernel_params * kparams = mctx->kparams;
@@ -742,21 +742,21 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
 
     uint32_t row_start = mctx->row_start;
     uint32_t row_end   = mctx->row_end;
-    uint32_t r_start   = row_start + ith * mctx->nrows_per_thread;
-    uint32_t r_end     = MIN(r_start + mctx->nrows_per_thread, row_end);
+    uint32_t r_start   = row_start + ith * mctx->nrows_per_hvx;
+    uint32_t r_end     = MIN(r_start + mctx->nrows_per_hvx, row_end);
 
     if (r_start >= r_end) return;
 
     size_t slot_size = mctx->phase1_slot_size;
     size_t spad_offset = ith * slot_size * (size_t) kparams->n_slots;
-    uint8_t * thread_spad = mctx->phase1_base + spad_offset;
+    uint8_t * hvx_spad = mctx->phase1_base + spad_offset;
 
     uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
-    float *   values_buf  = (float *) thread_spad;
-    int32_t * indices_buf = (int32_t *) (thread_spad + chunk_elems * sizeof(float));
+    float *   values_buf  = (float *) hvx_spad;
+    int32_t * indices_buf = (int32_t *) (hvx_spad + chunk_elems * sizeof(float));
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    dma_queue * dma_q = th->dma_q;
+    struct htp_thread_trace * tr = th->trace;
 
     bool is_top_k = kparams->is_top_k != 0;
     uint32_t k = dst->ne[0];
@@ -768,7 +768,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
     if (kparams->n_slots == 2) {
         struct htp_tensor_row_coord prime_coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
         for (uint32_t r = r_start, spad_idx = 0; r < r_end && spad_idx < 2; r++, spad_idx++) {
-            uint8_t * cur_spad = thread_spad + spad_idx * slot_size;
+            uint8_t * cur_spad = hvx_spad + spad_idx * slot_size;
             float *   cur_values  = (float *) cur_spad;
             int32_t * cur_indices = (int32_t *) (cur_spad + chunk_elems * sizeof(float));
 
@@ -791,7 +791,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
 
         for (uint32_t r = r_start; r < r_end; r++) {
             uint32_t cur_slot = (r - r_start) & 1;
-            uint8_t * cur_spad = thread_spad + cur_slot * slot_size;
+            uint8_t * cur_spad = hvx_spad + cur_slot * slot_size;
             float *   cur_values  = (float *) cur_spad;
             int32_t * cur_indices = (int32_t *) (cur_spad + chunk_elems * sizeof(float));
 
@@ -876,7 +876,7 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
 
     const struct htp_sort_kernel_params * kparams = (const struct htp_sort_kernel_params *) octx->kernel_params;
 
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+    if (!htp_ops_context_set_n_hvx(octx, kparams->n_hvx)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
@@ -886,7 +886,7 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                                (uint32_t) kparams->chunk_elems,
                                (uint32_t) kparams->n_chunks,
                                (uint32_t) kparams->merge_elems,
-                               (uint32_t) kparams->n_threads,
+                               (uint32_t) kparams->n_hvx,
                                (uint32_t) kparams->n_slots);
 
     if (octx->ctx->vtcm_size < layout.total_bytes) {
@@ -909,7 +909,7 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
         cctx.merge_indices    = merge_indices;
         memset(cctx.real_count, 0, sizeof(cctx.real_count));
 
-        work_queue_run(octx->ctx->work_queue, htp_sort_chunk_job, &cctx, (uint32_t) kparams->n_chunks);
+        work_queue_run(octx->ctx->hvx_queue, htp_sort_chunk_job, &cctx, (uint32_t) kparams->n_chunks);
 
         if (is_top_k) {
             uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
@@ -952,10 +952,10 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
             bool asc = kparams->order == GGML_SORT_ORDER_ASC;
             if (kparams->n_chunks == 8) {
                 struct htp_argsort_merge_l1_ctx l1ctx = { &cctx, asc };
-                work_queue_run(octx->ctx->work_queue, htp_argsort_merge_l1_job, &l1ctx, 4);
+                work_queue_run(octx->ctx->hvx_queue, htp_argsort_merge_l1_job, &l1ctx, 4);
 
                 struct htp_argsort_merge_l2_ctx l2ctx = { &cctx, asc };
-                work_queue_run(octx->ctx->work_queue, htp_argsort_merge_l2_job, &l2ctx, 2);
+                work_queue_run(octx->ctx->hvx_queue, htp_argsort_merge_l2_job, &l2ctx, 2);
 
                 size_t slot_size = layout.phase1_slot_size;
                 uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
@@ -976,7 +976,7 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                 dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
             } else if (kparams->n_chunks == 4) {
                 struct htp_argsort_merge_l1_ctx l1ctx = { &cctx, asc };
-                work_queue_run(octx->ctx->work_queue, htp_argsort_merge_l1_job, &l1ctx, 2);
+                work_queue_run(octx->ctx->hvx_queue, htp_argsort_merge_l1_job, &l1ctx, 2);
 
                 uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
                 uint32_t n01 = cctx.real_count[0] + cctx.real_count[1];
@@ -1031,13 +1031,13 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
     if (rows_per_chunk == 0) {
         rows_per_chunk = 1;
     }
-    uint32_t dr = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div);
+    uint32_t dr = fastdiv(nrows + octx->n_hvx - 1, &octx->n_hvx_div);
     if (rows_per_chunk > 1) {
         dr = hex_round_up(dr, rows_per_chunk);
     }
-    mctx.nrows_per_thread = dr;
+    mctx.nrows_per_hvx = dr;
 
-    work_queue_run(octx->ctx->work_queue, htp_sort_multi_row_job, &mctx, octx->n_threads);
+    work_queue_run(octx->ctx->hvx_queue, htp_sort_multi_row_job, &mctx, octx->n_hvx);
 
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/argsort-ops.h

@@ -10,7 +10,7 @@
 #include "htp-vtcm.h"
 
 struct htp_sort_kernel_params {
-    uint32_t n_threads;
+    uint32_t n_hvx;
     uint32_t order;             // GGML_SORT_ORDER_ASC (0) or GGML_SORT_ORDER_DESC (1)
     uint32_t is_top_k;          // 1 if TOP_K, 0 if ARGSORT
     uint32_t chunk_elems;
@@ -31,7 +31,7 @@ struct htp_sort_vtcm_layout {
     uint32_t chunk_elems;
     uint32_t n_chunks;
     uint32_t merge_elems;
-    uint32_t n_threads;
+    uint32_t n_hvx;
     uint32_t n_slots;
 };
 
@@ -41,14 +41,14 @@ static inline void htp_sort_vtcm_layout_build(
     uint32_t chunk_elems,
     uint32_t n_chunks,
     uint32_t merge_elems,
-    uint32_t n_threads,
+    uint32_t n_hvx,
     uint32_t n_slots) {
 
     memset(L, 0, sizeof(*L));
     L->chunk_elems = chunk_elems;
     L->n_chunks    = n_chunks;
     L->merge_elems = merge_elems;
-    L->n_threads   = n_threads;
+    L->n_hvx       = n_hvx;
     L->n_slots     = n_slots;
 
     const size_t values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
@@ -57,7 +57,7 @@ static inline void htp_sort_vtcm_layout_build(
 
     size_t off = 0;
     if (total_rows > 1) {
-        VTCM_LAYOUT_ALLOC(off, off_phase1, L->phase1_slot_size * n_slots * n_threads);
+        VTCM_LAYOUT_ALLOC(off, off_phase1, L->phase1_slot_size * n_slots * n_hvx);
     } else {
         VTCM_LAYOUT_ALLOC(off, off_phase1, L->phase1_slot_size * n_chunks * n_slots);
         const size_t merge_values_size  = ((merge_elems * sizeof(float)) + 127) & ~127;
@@ -73,14 +73,14 @@ static inline bool htp_sort_solve_layout(
     uint32_t ne00,
     uint32_t total_rows,
     uint32_t k,
-    uint32_t n_threads_max,
+    uint32_t n_hvx_max,
     size_t vtcm_budget,
     bool is_top_k) {
 
     memset(layout, 0, sizeof(*layout));
 
     if (total_rows > 1) {
-        uint32_t n_threads = total_rows < n_threads_max ? total_rows : n_threads_max;
+        uint32_t n_hvx = total_rows < n_hvx_max ? total_rows : n_hvx_max;
 
         uint32_t n_vec = (ne00 + 31) / 32;
         uint32_t n_vec_pow2 = 1;
@@ -96,13 +96,13 @@ static inline bool htp_sort_solve_layout(
             n_slots = 1;
         }
 
-        size_t spad_per_thread = spad_per_slot * n_slots;
+        size_t spad_per_hvx = spad_per_slot * n_slots;
 
-        while (n_threads > 1 && (spad_per_thread * n_threads) > vtcm_budget) {
-            n_threads--;
+        while (n_hvx > 1 && (spad_per_hvx * n_hvx) > vtcm_budget) {
+            n_hvx--;
         }
 
-        htp_sort_vtcm_layout_build(layout, total_rows, ne00_padded, 1, 0, n_threads, n_slots);
+        htp_sort_vtcm_layout_build(layout, total_rows, ne00_padded, 1, 0, n_hvx, n_slots);
         return layout->total_bytes <= vtcm_budget;
     }
 
@@ -112,7 +112,7 @@ static inline bool htp_sort_solve_layout(
 
     uint32_t n_chunks = 1;
     if (ne00 > 1024) {
-        while (n_chunks * 2 <= n_threads_max && n_chunks * 2 <= n_vec_pow2) {
+        while (n_chunks * 2 <= n_hvx_max && n_chunks * 2 <= n_vec_pow2) {
             n_chunks *= 2;
         }
     }

ggml/src/ggml-hexagon/htp/binary-ops.c

@@ -14,7 +14,7 @@
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
 #include "hex-common.h"
-#include "hex-profile.h"
+#include "htp-trace.h"
 #include "binary-ops.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
@@ -39,7 +39,7 @@ struct htp_binary_context {
     struct fastdiv_values src1_dim3_div; // ne13
 
     uint32_t block_max;
-    uint32_t nrows_per_thread;
+    uint32_t nrows_per_hvx;
     uint32_t total_rows;
     uint32_t row_start;
     uint32_t n_srcs;
@@ -260,7 +260,7 @@ static void compute_add_id_f32(
 }
 
 // 1. Scalar src1 in VTCM via DMA (ne10 == 1)
-static void binary_thread_scalar(unsigned int nth, unsigned int ith, void * data) {
+static void binary_thread_scalar(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct htp_binary_context * bctx = (struct htp_binary_context *) data;
     struct htp_ops_context * octx = bctx->octx;
     htp_binary_preamble;
@@ -270,21 +270,21 @@ static void binary_thread_scalar(unsigned int nth, unsigned int ith, void * data
     const size_t row_stride   = (size_t) ne11 * elem_size;
 
     const uint32_t row_size_bytes = bctx->row_size_bytes;
-    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
-    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
+    const uint32_t start_row = bctx->row_start + bctx->nrows_per_hvx * ith;
+    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_hvx, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
     FARF(HIGH, "binary-scalar: %d/%d (%u:%u) row-size %u (%u)",
          ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
 
     const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout;
-    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread);
-    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
+    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_hvx);
+    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_hvx);
     size_t src0_spad_half    = layout->src0_spad_half_size;
     size_t dst_spad_half     = layout->dst_spad_half_size;
     const void * s1_table    = VTCM_LAYOUT_PTR(const void, bctx->vtcm_base, layout->off_src1);
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     uint32_t ir_prefetch = start_row;
     int spad_idx = 0;
 
@@ -309,7 +309,7 @@ static void binary_thread_scalar(unsigned int nth, unsigned int ith, void * data
         spad_idx ^= 1;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
     compute_scalar_t compute = (compute_scalar_t) bctx->compute;
 
     for (uint32_t ir = start_row; ir < end_row; ) {
@@ -358,29 +358,29 @@ static void binary_thread_scalar(unsigned int nth, unsigned int ith, void * data
 }
 
 // 2. Vector Same Shape (ne1x == ne0x) or Simple Broadcast
-static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith, void * data) {
+static void binary_thread_vector_same_shape(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct htp_binary_context * bctx = (struct htp_binary_context *) data;
     struct htp_ops_context * octx = bctx->octx;
     htp_binary_preamble;
 
     const uint32_t row_size_bytes = bctx->row_size_bytes;
-    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
-    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
+    const uint32_t start_row = bctx->row_start + bctx->nrows_per_hvx * ith;
+    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_hvx, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
     FARF(HIGH, "binary-same-shape: %d/%d (%u:%u) row-size %u (%u)",
          ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
 
     const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout;
-    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread);
-    uint8_t * src1_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src1) + (ith * layout->src1_bytes_per_thread);
-    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
+    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_hvx);
+    uint8_t * src1_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src1) + (ith * layout->src1_bytes_per_hvx);
+    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_hvx);
 
     size_t src0_spad_half = layout->src0_spad_half_size;
     size_t src1_spad_half = layout->src1_spad_half_size;
     size_t dst_spad_half  = layout->dst_spad_half_size;
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     const uint32_t n_srcs = bctx->n_srcs;
     uint32_t ir_prefetch = start_row;
     int spad_idx = 0;
@@ -415,7 +415,7 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
         spad_idx ^= 1;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
     compute_same_shape_t compute = (compute_same_shape_t) bctx->compute;
 
     for (uint32_t ir = start_row; ir < end_row; ) {
@@ -472,27 +472,27 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
 }
 
 // 3. Row Broadcast (ne11 == 1, ne12 == 1, single row src1)
-static void binary_thread_vector_row_broadcast(unsigned int nth, unsigned int ith, void * data) {
+static void binary_thread_vector_row_broadcast(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct htp_binary_context * bctx = (struct htp_binary_context *) data;
     struct htp_ops_context * octx = bctx->octx;
     htp_binary_preamble;
 
     const uint32_t row_size_bytes = bctx->row_size_bytes;
-    const uint32_t start_row  = bctx->row_start + bctx->nrows_per_thread * ith;
-    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
+    const uint32_t start_row = bctx->row_start + bctx->nrows_per_hvx * ith;
+    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_hvx, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
     FARF(HIGH, "binary-row-bcast: %d/%d (%u:%u) row-size %u (%u)",
          ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
 
     const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout;
-    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread);
-    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
+    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_hvx);
+    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_hvx);
 
     size_t src0_spad_half = layout->src0_spad_half_size;
     size_t dst_spad_half  = layout->dst_spad_half_size;
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     uint32_t ir_prefetch = start_row;
     int spad_idx = 0;
 
@@ -517,7 +517,7 @@ static void binary_thread_vector_row_broadcast(unsigned int nth, unsigned int it
         spad_idx ^= 1;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
     compute_row_bcast_t compute = (compute_row_bcast_t) bctx->compute;
 
     for (uint32_t ir = start_row; ir < end_row; ) {
@@ -555,22 +555,22 @@ static void binary_thread_vector_row_broadcast(unsigned int nth, unsigned int it
 }
 
 // 4. Vector Complex (ne10 == ne00, complex broadcast)
-static void binary_thread_vector_complex(unsigned int nth, unsigned int ith, void * data) {
+static void binary_thread_vector_complex(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct htp_binary_context * bctx = (struct htp_binary_context *) data;
     struct htp_ops_context * octx = bctx->octx;
     htp_binary_preamble;
 
     const uint32_t row_size_bytes = bctx->row_size_bytes;
-    const uint32_t start_row  = bctx->row_start + bctx->nrows_per_thread * ith;
-    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
+    const uint32_t start_row = bctx->row_start + bctx->nrows_per_hvx * ith;
+    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_hvx, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
     FARF(HIGH, "binary-complex: %d/%d (%u:%u) row-size %u (%u)",
          ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
 
     const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout;
-    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread);
-    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
+    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_hvx);
+    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_hvx);
     size_t src0_spad_half    = layout->src0_spad_half_size;
     size_t dst_spad_half     = layout->dst_spad_half_size;
     const uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(const uint8_t, bctx->vtcm_base, layout->off_src1);
@@ -580,7 +580,7 @@ static void binary_thread_vector_complex(unsigned int nth, unsigned int ith, voi
     const size_t vtcm_plane_stride = (size_t) ne11 * vtcm_row_stride;
     const size_t vtcm_vol_stride   = (size_t) ne12 * vtcm_plane_stride;
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     uint32_t ir_prefetch = start_row;
     int spad_idx = 0;
 
@@ -603,7 +603,7 @@ static void binary_thread_vector_complex(unsigned int nth, unsigned int ith, voi
         spad_idx ^= 1;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
     compute_complex_t compute = (compute_complex_t) bctx->compute;
 
     for (uint32_t ir = start_row; ir < end_row; ) {
@@ -646,19 +646,19 @@ static void binary_thread_vector_complex(unsigned int nth, unsigned int ith, voi
 }
 
 // 5. Element Repeat (ne10 != ne00)
-static void binary_thread_element_repeat(unsigned int nth, unsigned int ith, void * data) {
+static void binary_thread_element_repeat(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct htp_binary_context * bctx = (struct htp_binary_context *) data;
     struct htp_ops_context * octx = bctx->octx;
     htp_binary_preamble;
 
     const uint32_t row_size_bytes = bctx->row_size_bytes;
-    const uint32_t start_row  = bctx->row_start + bctx->nrows_per_thread * ith;
-    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
+    const uint32_t start_row = bctx->row_start + bctx->nrows_per_hvx * ith;
+    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_hvx, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
     const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout;
-    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread);
-    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
+    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_hvx);
+    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_hvx);
     size_t src0_spad_half    = layout->src0_spad_half_size;
     size_t dst_spad_half     = layout->dst_spad_half_size;
     const uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(const uint8_t, bctx->vtcm_base, layout->off_src1);
@@ -671,7 +671,7 @@ static void binary_thread_element_repeat(unsigned int nth, unsigned int ith, voi
     FARF(HIGH, "binary-repeat: %d/%d (%u:%u) row-size %u (%u)",
          ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     uint32_t ir_prefetch = start_row;
     int spad_idx = 0;
 
@@ -694,7 +694,7 @@ static void binary_thread_element_repeat(unsigned int nth, unsigned int ith, voi
         spad_idx ^= 1;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
     compute_repeat_t compute = (compute_repeat_t) bctx->compute;
 
     for (uint32_t ir = start_row; ir < end_row; ) {
@@ -737,7 +737,7 @@ static void binary_thread_element_repeat(unsigned int nth, unsigned int ith, voi
 }
 
 // 6. ADD_ID (src1 gathered via src2 indices)
-static void binary_thread_add_id_f32(unsigned int nth, unsigned int ith, void * data) {
+static void binary_thread_add_id_f32(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct htp_binary_context * bctx = (struct htp_binary_context *) data;
     struct htp_ops_context * octx = bctx->octx;
 
@@ -760,19 +760,19 @@ static void binary_thread_add_id_f32(unsigned int nth, unsigned int ith, void *
     const uint32_t nb3 = dst->nb[3];
 
     const uint32_t row_size_bytes = bctx->row_size_bytes;
-    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
-    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
+    const uint32_t start_row = bctx->row_start + bctx->nrows_per_hvx * ith;
+    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_hvx, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
     const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout;
-    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread);
-    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
+    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_hvx);
+    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_hvx);
     const uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(const uint8_t, bctx->vtcm_base, layout->off_src1);
     const int32_t * vtcm_src2 = VTCM_LAYOUT_PTR(const int32_t, bctx->vtcm_base, layout->off_src2);
     size_t src0_spad_half    = layout->src0_spad_half_size;
     size_t dst_spad_half     = layout->dst_spad_half_size;
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     uint32_t ir_prefetch = start_row;
     int spad_idx = 0;
 
@@ -795,7 +795,7 @@ static void binary_thread_add_id_f32(unsigned int nth, unsigned int ith, void *
         spad_idx ^= 1;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     for (uint32_t ir = start_row; ir < end_row; ) {
         uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
@@ -870,7 +870,7 @@ struct binary_chunked_context {
     uint32_t                 elem_start;
     uint32_t                 nelem;
     uint32_t                 chunk_size;
-    uint32_t                 chunks_per_thread;
+    uint32_t                 chunks_per_hvx;
     uint32_t                 total_chunks;
     uint32_t                 n_srcs;
     bool                     is_scalar;
@@ -881,15 +881,15 @@ struct binary_chunked_context {
     compute_binary_scalar_chunked_f16_t compute_scalar_f16;
 };
 
-static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * data) {
+static void binary_thread_chunked(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     (void) nth;
     struct binary_chunked_context * ctx = (struct binary_chunked_context *) data;
     struct htp_ops_context * octx = ctx->octx;
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
 
-    const uint32_t start_chunk = ctx->chunks_per_thread * ith;
-    const uint32_t end_chunk   = MIN(start_chunk + ctx->chunks_per_thread, ctx->total_chunks);
+    const uint32_t start_chunk = ctx->chunks_per_hvx * ith;
+    const uint32_t end_chunk   = MIN(start_chunk + ctx->chunks_per_hvx, ctx->total_chunks);
     if (start_chunk >= end_chunk) {
         return;
     }
@@ -903,11 +903,11 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
          ith, nth, start_chunk, end_chunk, ctx->total_chunks, ctx->nelem);
 
     const struct htp_binary_vtcm_layout * layout = &ctx->vtcm_layout;
-    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread);
-    uint8_t * src1_spad_base = ctx->is_scalar ? NULL : (VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_src1) + (ith * layout->src1_bytes_per_thread));
-    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
+    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_hvx);
+    uint8_t * src1_spad_base = ctx->is_scalar ? NULL : (VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_src1) + (ith * layout->src1_bytes_per_hvx));
+    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_hvx);
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     const uint32_t n_srcs = ctx->n_srcs;
     uint32_t prefetch_chunk = start_chunk;
     int spad_idx = 0;
@@ -939,7 +939,7 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
         spad_idx ^= 1;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     for (uint32_t c = start_chunk; c < end_chunk; c++) {
         const uint32_t c_start   = ctx->elem_start + c * chunk_size;
@@ -1091,7 +1091,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
             return HTP_STATUS_OK;
         }
 
-        if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        if (!htp_ops_context_set_n_hvx(octx, kparams->n_hvx)) {
             return HTP_STATUS_INVAL_PARAMS;
         }
 
@@ -1103,8 +1103,8 @@ static int execute_op_binary(struct htp_ops_context * octx) {
 
         const uint32_t chunk_size = kparams->chunk_size;
         const uint32_t total_chunks = (nelem + chunk_size - 1) / chunk_size;
-        const uint32_t n_threads = (total_chunks >= 2) ? MIN(octx->n_threads, total_chunks) : 1;
-        const uint32_t chunks_per_thread = (total_chunks + n_threads - 1) / n_threads;
+        const uint32_t n_hvx = (total_chunks >= 2) ? MIN(octx->n_hvx, total_chunks) : 1;
+        const uint32_t chunks_per_hvx = (total_chunks + n_hvx - 1) / n_hvx;
 
         float scalar_f32 = 0.0f;
         _Float16 scalar_f16 = 0;
@@ -1128,7 +1128,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
             .elem_start          = elem_start,
             .nelem               = nelem,
             .chunk_size          = chunk_size,
-            .chunks_per_thread   = chunks_per_thread,
+            .chunks_per_hvx      = chunks_per_hvx,
             .total_chunks        = total_chunks,
             .n_srcs              = n_srcs,
             .is_scalar           = is_scalar,
@@ -1139,7 +1139,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
             .compute_scalar_f16  = compute_scalar_f16,
         };
 
-        work_queue_run(octx->ctx->work_queue, binary_thread_chunked, &cctx, n_threads);
+        work_queue_run(octx->ctx->hvx_queue, binary_thread_chunked, &cctx, n_hvx);
         return HTP_STATUS_OK;
     }
 
@@ -1159,11 +1159,11 @@ static int execute_op_binary(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+    if (!htp_ops_context_set_n_hvx(octx, kparams->n_hvx)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_hvx = octx->n_hvx;
     const size_t src0_row_size_aligned = kparams->src0_row_size_aligned;
     const size_t src1_row_size_aligned = kparams->src1_row_size_aligned;
     const size_t dst_row_size_aligned  = kparams->dst_row_size_aligned;
@@ -1176,7 +1176,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    worker_callback_t worker_func = NULL;
+    work_queue_func_t worker_func = NULL;
     void * compute_func = NULL;
 
     switch (kparams->kernel_type) {
@@ -1336,7 +1336,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
     }
 
     bctx.octx                  = octx;
-    bctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    bctx.nrows_per_hvx         = fastdiv(nrows + n_hvx - 1, &octx->n_hvx_div);
     bctx.total_rows            = nrows;
     bctx.row_start             = row_start;
     bctx.n_srcs                = n_srcs;
@@ -1375,7 +1375,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
         }
     }
 
-    work_queue_run(octx->ctx->work_queue, worker_func, &bctx, n_threads);
+    work_queue_run(octx->ctx->hvx_queue, worker_func, &bctx, n_hvx);
 
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/binary-ops.h

@@ -30,7 +30,7 @@ struct htp_binary_kernel_params {
     uint8_t  is_scalar;
     uint8_t  pad[2];
 
-    uint32_t n_threads;
+    uint32_t n_hvx;
     uint32_t rows_per_buffer;
 
     uint32_t src0_row_size_aligned;
@@ -86,9 +86,9 @@ struct htp_binary_vtcm_layout {
     size_t off_src2;
     size_t off_dst;
 
-    size_t src0_bytes_per_thread;
-    size_t src1_bytes_per_thread;
-    size_t dst_bytes_per_thread;
+    size_t src0_bytes_per_hvx;
+    size_t src1_bytes_per_hvx;
+    size_t dst_bytes_per_hvx;
 
     size_t src0_spad_half_size;
     size_t src1_spad_half_size;
@@ -106,8 +106,8 @@ static inline void htp_binary_vtcm_layout_build(
 ) {
     memset(L, 0, sizeof(*L));
 
-    const uint32_t n_threads = kparams->n_threads;
-    if (n_threads == 0) {
+    const uint32_t n_hvx = kparams->n_hvx;
+    if (n_hvx == 0) {
         return;
     }
 
@@ -119,9 +119,9 @@ static inline void htp_binary_vtcm_layout_build(
             return;
         }
 
-        L->src0_bytes_per_thread = 2 * chunk_bytes;
-        L->src1_bytes_per_thread = kparams->is_scalar ? 0 : (2 * (n_srcs - 1) * chunk_bytes);
-        L->dst_bytes_per_thread  = 2 * chunk_bytes;
+        L->src0_bytes_per_hvx = 2 * chunk_bytes;
+        L->src1_bytes_per_hvx = kparams->is_scalar ? 0 : (2 * (n_srcs - 1) * chunk_bytes);
+        L->dst_bytes_per_hvx  = 2 * chunk_bytes;
 
         L->src0_spad_half_size = chunk_bytes;
         L->src1_spad_half_size = kparams->is_scalar ? 0 : chunk_bytes;
@@ -130,9 +130,9 @@ static inline void htp_binary_vtcm_layout_build(
         L->rows_per_buffer = 1;
         L->src1_size = 0;
 
-        const size_t src0_total = n_threads * L->src0_bytes_per_thread;
-        const size_t src1_total = kparams->is_scalar ? 128 : (n_threads * L->src1_bytes_per_thread);
-        const size_t dst_total  = n_threads * L->dst_bytes_per_thread;
+        const size_t src0_total = n_hvx * L->src0_bytes_per_hvx;
+        const size_t src1_total = kparams->is_scalar ? 128 : (n_hvx * L->src1_bytes_per_hvx);
+        const size_t dst_total  = n_hvx * L->dst_bytes_per_hvx;
 
         size_t off = 0;
         VTCM_LAYOUT_ALLOC(off, off_src0, src0_total);
@@ -158,7 +158,7 @@ static inline void htp_binary_vtcm_layout_build(
 
     const size_t rows_per_buffer = (kparams->rows_per_buffer > 0)
         ? kparams->rows_per_buffer
-        : (vtcm_size - static_vtcm_size) / (n_threads * spad_row_total);
+        : (vtcm_size - static_vtcm_size) / (n_hvx * spad_row_total);
     if (rows_per_buffer == 0) {
         return;
     }
@@ -167,24 +167,24 @@ static inline void htp_binary_vtcm_layout_build(
     L->src1_size = kparams->src1_size;
     L->src2_size = kparams->src2_size;
 
-    L->src0_bytes_per_thread = rows_per_buffer * 2 * kparams->src0_row_size_aligned;
-    L->dst_bytes_per_thread  = rows_per_buffer * 2 * kparams->dst_row_size_aligned;
-    L->src1_bytes_per_thread = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
+    L->src0_bytes_per_hvx = rows_per_buffer * 2 * kparams->src0_row_size_aligned;
+    L->dst_bytes_per_hvx  = rows_per_buffer * 2 * kparams->dst_row_size_aligned;
+    L->src1_bytes_per_hvx = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
         ? rows_per_buffer * 2 * (n_srcs - 1) * kparams->src1_row_size_aligned
         : 0;
 
-    L->src0_spad_half_size = L->src0_bytes_per_thread / 2;
+    L->src0_spad_half_size = L->src0_bytes_per_hvx / 2;
     L->src1_spad_half_size = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
         ? rows_per_buffer * kparams->src1_row_size_aligned
         : 0;
-    L->dst_spad_half_size  = L->dst_bytes_per_thread / 2;
+    L->dst_spad_half_size  = L->dst_bytes_per_hvx / 2;
 
-    const size_t src0_total = n_threads * L->src0_bytes_per_thread;
+    const size_t src0_total = n_hvx * L->src0_bytes_per_hvx;
     const size_t src1_total = (kparams->src1_size > 0)
         ? kparams->src1_size
-        : n_threads * L->src1_bytes_per_thread;
+        : n_hvx * L->src1_bytes_per_hvx;
     const size_t src2_total = kparams->src2_size;
-    const size_t dst_total  = n_threads * L->dst_bytes_per_thread;
+    const size_t dst_total  = n_hvx * L->dst_bytes_per_hvx;
 
     size_t off = 0;
     VTCM_LAYOUT_ALLOC(off, off_src0, src0_total);
@@ -200,15 +200,15 @@ static inline bool htp_binary_solve_chunked(
     struct htp_binary_kernel_params * kparams,
     uint32_t total_elems,
     uint32_t elem_size,
-    uint32_t max_threads,
+    uint32_t max_hvx,
     uint32_t n_srcs,
     bool is_scalar,
     size_t vtcm_size
 ) {
     const size_t n_bufs = is_scalar ? 2 : (size_t) n_srcs + 1;
     const size_t scalar_bytes = is_scalar ? 128 : 0;
-    const size_t budget_per_buf = (max_threads > 0 && vtcm_size > scalar_bytes)
-        ? hex_align_down((vtcm_size - scalar_bytes) / (2 * max_threads * n_bufs), 128)
+    const size_t budget_per_buf = (max_hvx > 0 && vtcm_size > scalar_bytes)
+        ? hex_align_down((vtcm_size - scalar_bytes) / (2 * max_hvx * n_bufs), 128)
         : 0;
     const uint32_t max_chunk_bytes = (budget_per_buf > 0 && budget_per_buf < 32768) ? (uint32_t) budget_per_buf : 32768u;
     const uint32_t max_chunk_elems = max_chunk_bytes / elem_size;
@@ -217,16 +217,16 @@ static inline bool htp_binary_solve_chunked(
     }
 
     const uint32_t min_chunk_elems = 256 < max_chunk_elems ? 256 : max_chunk_elems;
-    const uint32_t target_chunk_elems = hex_round_up((total_elems + (2 * max_threads) - 1) / (2 * max_threads), 32);
+    const uint32_t target_chunk_elems = hex_round_up((total_elems + (2 * max_hvx) - 1) / (2 * max_hvx), 32);
     const uint32_t chunk_size  = (target_chunk_elems < min_chunk_elems) ? min_chunk_elems
                                : (target_chunk_elems > max_chunk_elems) ? max_chunk_elems
                                : target_chunk_elems;
     const uint32_t chunk_bytes  = hex_round_up(chunk_size * elem_size, 128);
     const uint32_t total_chunks = (total_elems + chunk_size - 1) / chunk_size;
-    const uint32_t n_threads    = (total_chunks >= 2) ? (max_threads < total_chunks ? max_threads : total_chunks) : 1;
+    const uint32_t n_hvx        = (total_chunks >= 2) ? (max_hvx < total_chunks ? max_hvx : total_chunks) : 1;
 
     kparams->kernel_type           = HTP_BINARY_KERNEL_CHUNKED;
-    kparams->n_threads             = n_threads;
+    kparams->n_hvx                 = n_hvx;
     kparams->n_srcs                = (uint8_t) n_srcs;
     kparams->rows_per_buffer       = 1;
     kparams->src1_row_size_aligned = is_scalar ? 0 : kparams->src0_row_size_aligned;

ggml/src/ggml-hexagon/htp/concat-ops.c

@@ -3,7 +3,7 @@
 #include "hex-common.h"
 #include "dma-copy.h"
 #include "hex-fastdiv.h"
-#include "hex-profile.h"
+#include "htp-trace.h"
 #include "hexagon_protos.h"
 #include "hexagon_types.h"
 #include "htp-ctx.h"
@@ -21,11 +21,11 @@ struct htp_concat_context {
     struct htp_ops_context * octx;
     uint8_t * spad0_base;
     uint8_t * spad1_base;
-    uint32_t  spad0_size_per_thread;
-    uint32_t  spad1_size_per_thread;
+    uint32_t  spad0_size_per_hvx;
+    uint32_t  spad1_size_per_hvx;
     uint32_t  row_start;
     uint32_t  nrows;
-    uint32_t  nrows_per_thread;
+    uint32_t  nrows_per_hvx;
     uint32_t  nplanes;
     struct fastdiv_values div_ne2;
 };
@@ -36,7 +36,7 @@ static inline dma_addr_t concat_plane_addr(const struct htp_tensor * t, uint32_t
     return t->data + i2 * t->nb[2] + i3 * t->nb[3];
 }
 
-static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void * data) {
+static void concat_2d_f32_transposed(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct htp_concat_context * cctx = (struct htp_concat_context *) data;
     struct htp_ops_context * octx = cctx->octx;
 
@@ -48,12 +48,12 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
     const uint32_t src1_ne0 = src1->ne[0];
 
     // dst rows [row_begin, row_end) of this thread; a dst row is a src0 row followed by a src1 column
-    const uint32_t row_begin = cctx->row_start + ith * cctx->nrows_per_thread;
-    const uint32_t row_end   = MIN(row_begin + cctx->nrows_per_thread, cctx->row_start + cctx->nrows);
+    const uint32_t row_begin = cctx->row_start + ith * cctx->nrows_per_hvx;
+    const uint32_t row_end   = MIN(row_begin + cctx->nrows_per_hvx, cctx->row_start + cctx->nrows);
     if (row_begin >= row_end) return;
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    dma_queue * dma_q = th->dma_q;
+    struct htp_thread_trace * tr = th->trace;
 
     const struct fastdiv_values * div_ne2 = &cctx->div_ne2;
     const uint32_t ne2 = dst->ne[2];
@@ -72,8 +72,8 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
     const uint32_t spad0_stride      = src0_row_padded + block_cols_padded * sizeof(float);
 
     // two slots per scratchpad: tile b gathers in slot b&1 while dst[b-1] and src1[b+1] DMA in the other
-    uint8_t * spad0_base = cctx->spad0_base + ith * cctx->spad0_size_per_thread;
-    uint8_t * spad1_base = cctx->spad1_base + ith * cctx->spad1_size_per_thread;
+    uint8_t * spad0_base = cctx->spad0_base + ith * cctx->spad0_size_per_hvx;
+    uint8_t * spad1_base = cctx->spad1_base + ith * cctx->spad1_size_per_hvx;
     const uint32_t spad0_slot = block_rows * spad0_stride;
     const uint32_t spad1_slot = block_cols_padded * spad1_stride;
 
@@ -153,7 +153,7 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
     dma_queue_flush(dma_q);
 }
 
-static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void * data) {
+static void concat_2d_f16_transposed(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct htp_concat_context * cctx = (struct htp_concat_context *) data;
     struct htp_ops_context * octx = cctx->octx;
 
@@ -165,12 +165,12 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
     const uint32_t src1_ne0 = src1->ne[0];
 
     // dst rows [row_begin, row_end) of this thread; a dst row is a src0 row followed by a src1 column
-    const uint32_t row_begin = cctx->row_start + ith * cctx->nrows_per_thread;
-    const uint32_t row_end   = MIN(row_begin + cctx->nrows_per_thread, cctx->row_start + cctx->nrows);
+    const uint32_t row_begin = cctx->row_start + ith * cctx->nrows_per_hvx;
+    const uint32_t row_end   = MIN(row_begin + cctx->nrows_per_hvx, cctx->row_start + cctx->nrows);
     if (row_begin >= row_end) return;
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    dma_queue * dma_q = th->dma_q;
+    struct htp_thread_trace * tr = th->trace;
 
     const struct fastdiv_values * div_ne2 = &cctx->div_ne2;
     const uint32_t ne2 = dst->ne[2];
@@ -189,8 +189,8 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
     const uint32_t spad0_stride      = src0_row_padded + block_cols_padded * sizeof(__fp16);
 
     // two slots per scratchpad: tile b gathers in slot b&1 while dst[b-1] and src1[b+1] DMA in the other
-    uint8_t * spad0_base = cctx->spad0_base + ith * cctx->spad0_size_per_thread;
-    uint8_t * spad1_base = cctx->spad1_base + ith * cctx->spad1_size_per_thread;
+    uint8_t * spad0_base = cctx->spad0_base + ith * cctx->spad0_size_per_hvx;
+    uint8_t * spad1_base = cctx->spad1_base + ith * cctx->spad1_size_per_hvx;
     const uint32_t spad0_slot = block_rows * spad0_stride;
     const uint32_t spad1_slot = block_cols_padded * spad1_stride;
 
@@ -311,7 +311,7 @@ static int concat_regular(struct htp_ops_context * octx, int dim, uint32_t type_
 }
 
 static int concat_transposed(struct htp_ops_context * octx, const struct htp_concat_kernel_params * kparams, uint32_t type_size) {
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+    if (!htp_ops_context_set_n_hvx(octx, kparams->n_hvx)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
@@ -334,10 +334,10 @@ static int concat_transposed(struct htp_ops_context * octx, const struct htp_con
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_hvx = octx->n_hvx;
 
     struct htp_concat_transposed_vtcm_layout layout;
-    htp_concat_transposed_vtcm_layout_build(&layout, octx->src[0]->ne[0], octx->src[1]->ne[0], type_size, n_threads);
+    htp_concat_transposed_vtcm_layout_build(&layout, octx->src[0]->ne[0], octx->src[1]->ne[0], type_size, n_hvx);
 
     uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
 
@@ -345,16 +345,16 @@ static int concat_transposed(struct htp_ops_context * octx, const struct htp_con
     cctx.octx                  = octx;
     cctx.spad0_base            = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_spad0);
     cctx.spad1_base            = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_spad1);
-    cctx.spad0_size_per_thread = layout.src0_spad_size_per_thread;
-    cctx.spad1_size_per_thread = layout.src1_spad_size_per_thread;
+    cctx.spad0_size_per_hvx    = layout.src0_spad_size_per_hvx;
+    cctx.spad1_size_per_hvx    = layout.src1_spad_size_per_hvx;
     cctx.row_start             = row_start;
     cctx.nrows                 = nrows;
     cctx.nplanes               = dst->ne[2] * dst->ne[3];
     cctx.div_ne2               = kparams->u.transposed.div_ne2;
-    cctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    cctx.nrows_per_hvx         = fastdiv(nrows + n_hvx - 1, &octx->n_hvx_div);
 
     work_queue_func_t worker_func = (type_size == 4) ? concat_2d_f32_transposed : concat_2d_f16_transposed;
-    work_queue_run(octx->ctx->work_queue, worker_func, &cctx, n_threads);
+    work_queue_run(octx->ctx->hvx_queue, worker_func, &cctx, n_hvx);
     return HTP_STATUS_OK;
 }
 
@@ -377,7 +377,7 @@ int op_concat(struct htp_ops_context * octx) {
             break;
 
         case HTP_CONCAT_KERNEL_GATHER_ROWS:
-            status = htp_gather_rows(octx, src0, src1, &kparams->u.gather, kparams->n_threads);
+            status = htp_gather_rows(octx, src0, src1, &kparams->u.gather, kparams->n_hvx);
             break;
 
         default:

ggml/src/ggml-hexagon/htp/concat-ops.h

@@ -21,7 +21,7 @@ struct htp_concat_transposed_params {
 struct htp_concat_kernel_params {
     uint8_t  kernel_type;
     uint8_t  dim;
-    uint8_t  n_threads;
+    uint8_t  n_hvx;
     uint8_t  pad;
 
     uint32_t vtcm_size;
@@ -39,8 +39,8 @@ _Static_assert(sizeof(struct htp_concat_kernel_params) <= 128, "htp_concat_kerne
 #endif
 
 struct htp_concat_transposed_vtcm_layout {
-    uint32_t src0_spad_size_per_thread;
-    uint32_t src1_spad_size_per_thread;
+    uint32_t src0_spad_size_per_hvx;
+    uint32_t src1_spad_size_per_hvx;
     size_t   off_spad0;
     size_t   off_spad1;
     uint32_t total_bytes;
@@ -51,7 +51,7 @@ static inline void htp_concat_transposed_vtcm_layout_build(
     uint32_t src0_ne0,
     uint32_t src1_ne0,
     uint32_t type_size,
-    uint32_t n_threads) {
+    uint32_t n_hvx) {
 
     // src1_ne0 is processed in chunks of block_cols tokens, so VTCM does not grow with the batch
     uint32_t block_rows = 128;
@@ -62,12 +62,12 @@ static inline void htp_concat_transposed_vtcm_layout_build(
     uint32_t spad0_row_bytes = hex_round_up(src0_ne0 * type_size, 128) + block_cols_padded * type_size;
 
     // double-buffered: two slots per thread
-    L->src0_spad_size_per_thread = 2 * block_rows * spad0_row_bytes;
-    L->src1_spad_size_per_thread = 2 * block_cols_padded * spad1_stride;
+    L->src0_spad_size_per_hvx = 2 * block_rows * spad0_row_bytes;
+    L->src1_spad_size_per_hvx = 2 * block_cols_padded * spad1_stride;
 
     size_t off = 0;
-    VTCM_LAYOUT_ALLOC(off, off_spad0, (size_t) n_threads * L->src0_spad_size_per_thread);
-    VTCM_LAYOUT_ALLOC(off, off_spad1, (size_t) n_threads * L->src1_spad_size_per_thread);
+    VTCM_LAYOUT_ALLOC(off, off_spad0, (size_t) n_hvx * L->src0_spad_size_per_hvx);
+    VTCM_LAYOUT_ALLOC(off, off_spad1, (size_t) n_hvx * L->src1_spad_size_per_hvx);
     L->total_bytes = (uint32_t) off;
 }
 

ggml/src/ggml-hexagon/htp/conv2d-ops.c

@@ -12,7 +12,7 @@
 #include "dma-queue.h"
 #include "hex-common.h"
 #include "hex-fastdiv.h"
-#include "hex-profile.h"
+#include "htp-trace.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
@@ -100,7 +100,7 @@ static void conv2d_mdev_rows(const struct htp_ops_context * octx, uint32_t total
     }
 }
 
-static void conv2d_hmx_worker(void * data) {
+static void conv2d_hmx_worker(const htp_thread_t * th, void * data) {
     struct htp_conv2d_hmx_task * task = (struct htp_conv2d_hmx_task *) data;
 
     const size_t dot_stride = task->n_dot_tiles * HMX_FP16_TILE_N_ELMS;
@@ -205,7 +205,7 @@ static inline uint32_t conv2d_hmx_chunk_range(const struct htp_conv2d_hmx_contex
     return cs;
 }
 
-static void conv2d_hmx_wpack_thread(unsigned int nth, unsigned int ith, void * data) {
+static void conv2d_hmx_wpack_thread(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     const struct htp_conv2d_hmx_context * hc = (const struct htp_conv2d_hmx_context *) data;
     struct htp_ops_context *                octx = hc->octx;
     const struct htp_conv2d_kernel_params * kp   = hc->kp;
@@ -218,7 +218,7 @@ static void conv2d_hmx_wpack_thread(unsigned int nth, unsigned int ith, void * d
     const bool     is_f16 = kernel->type == HTP_TYPE_F16;
     const uint32_t es     = is_f16 ? sizeof(__fp16) : sizeof(float);
 
-    dma_queue * dmaq = octx->ctx->dma[ith];
+    dma_queue * dmaq = th->dma_q;
     uint8_t *   thr  = hc->thr_base + (size_t) ith * kp->thr_size;
     const size_t w_buf_size = (size_t) WK * 32 * sizeof(__fp16);
     const size_t w_raw_size = (size_t) WK * 32 * sizeof(float);
@@ -232,7 +232,7 @@ static void conv2d_hmx_wpack_thread(unsigned int nth, unsigned int ith, void * d
 
     const HVX_Vector zero = Q6_V_vzero();
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     for (uint32_t rt = ith; rt < hc->n_rt; rt += nth) {
         const uint32_t co0 = (hc->rt0 + rt) * 32;
@@ -704,7 +704,7 @@ static inline bool conv2d_wpack_kmajor_next_task(struct conv2d_wpack_kmajor_iter
 }
 
 // Kernel stored as [k][cout] (cout contiguous): DMA k rows of 64 cout, transpose into two row tiles
-static void conv2d_hmx_wpack_kmajor_thread(unsigned int nth, unsigned int ith, void * data) {
+static void conv2d_hmx_wpack_kmajor_thread(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     const struct htp_conv2d_hmx_context * hc = (const struct htp_conv2d_hmx_context *) data;
     struct htp_ops_context *                octx = hc->octx;
     const struct htp_conv2d_kernel_params * kp   = hc->kp;
@@ -721,7 +721,7 @@ static void conv2d_hmx_wpack_kmajor_thread(unsigned int nth, unsigned int ith, v
     const uint32_t es     = is_f16 ? sizeof(__fp16) : sizeof(float);
     const uint32_t n_cb   = (hc->n_rt + 1) / 2;
 
-    dma_queue *  dmaq       = octx->ctx->dma[ith];
+    dma_queue *  dmaq       = th->dma_q;
     const size_t w_raw_size = (size_t) WK * 64 * sizeof(float);
     uint8_t *    wraws[2]   = { hc->thr_base + (size_t) ith * kp->thr_size,
                                 hc->thr_base + (size_t) ith * kp->thr_size + w_raw_size };
@@ -729,7 +729,7 @@ static void conv2d_hmx_wpack_kmajor_thread(unsigned int nth, unsigned int ith, v
     HVX_Vector * scratch    = vin + 32;
     HVX_Vector * sink       = scratch + 32;
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     struct conv2d_wpack_kmajor_iter it = {
         .cb    = 0,
@@ -778,7 +778,7 @@ static void conv2d_hmx_wpack_kmajor_thread(unsigned int nth, unsigned int ith, v
     }
 }
 
-static void conv2d_hmx_pack(const struct htp_conv2d_hmx_context * hc, uint32_t ith) {
+static void conv2d_hmx_pack(const struct htp_conv2d_hmx_context * hc, const htp_thread_t * th, uint32_t ith) {
     struct htp_ops_context *                octx = hc->octx;
     const struct htp_conv2d_kernel_params * kp   = hc->kp;
     const struct htp_tensor *               kernel = octx->src[0];
@@ -790,7 +790,7 @@ static void conv2d_hmx_pack(const struct htp_conv2d_hmx_context * hc, uint32_t i
     const uint32_t d0    = (uint32_t) hc->g->d0;
     const uint32_t n_vec = kp->tiled.n_k_tiles * (HMX_FP16_TILE_N_ELMS / 64);
 
-    dma_queue * dmaq  = octx->ctx->dma[ith];
+    dma_queue * dmaq  = th->dma_q;
     uint8_t *   thr   = hc->thr_base + (size_t) ith * kp->thr_size;
     float *     stages[2] = {
         (float *) thr,
@@ -808,7 +808,7 @@ static void conv2d_hmx_pack(const struct htp_conv2d_hmx_context * hc, uint32_t i
     uint32_t       t0, t1;
     const uint32_t cs = conv2d_hmx_chunk_range(hc, hc->pack_chunk, hc->pack_tpt, ith, &t0, &t1);
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     if (t0 >= t1) {
         return;
@@ -866,7 +866,7 @@ static void conv2d_hmx_pack(const struct htp_conv2d_hmx_context * hc, uint32_t i
     }
 }
 
-static void conv2d_hmx_unpack(const struct htp_conv2d_hmx_context * hc, uint32_t ith) {
+static void conv2d_hmx_unpack(const struct htp_conv2d_hmx_context * hc, const htp_thread_t * th, uint32_t ith) {
     struct htp_ops_context *                octx = hc->octx;
     const struct htp_conv2d_kernel_params * kp   = hc->kp;
     const struct htp_tensor *               dst  = octx->dst;
@@ -877,7 +877,7 @@ static void conv2d_hmx_unpack(const struct htp_conv2d_hmx_context * hc, uint32_t
     const uint32_t n_owt = kp->tiled.n_ow_tiles;
     const uint32_t OS    = kp->tiled.seg_tiles * 32;
 
-    dma_queue * dmaq      = octx->ctx->dma[ith];
+    dma_queue * dmaq      = th->dma_q;
     uint8_t *   thr       = hc->thr_base + (size_t) ith * kp->thr_size + (size_t) (kp->tiled.n_stage > 1 ? 2 : 1) * kp->stage_size;
     float *     ostage[2] = {
         (float *) thr,
@@ -889,7 +889,7 @@ static void conv2d_hmx_unpack(const struct htp_conv2d_hmx_context * hc, uint32_t
     const uint32_t cs  = conv2d_hmx_chunk_range(hc, hc->unpack_chunk, hc->unpack_tpt, ith, &t0, &t1);
     const uint32_t nct = MIN(cs + kp->ct_chunk, hc->t_end) - cs;
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     if (t0 >= t1) {
         return;
@@ -945,11 +945,11 @@ static void conv2d_hmx_unpack(const struct htp_conv2d_hmx_context * hc, uint32_t
     }
 }
 
-static void conv2d_hmx_phase_thread(unsigned int nth, unsigned int ith, void * data) {
+static void conv2d_hmx_phase_thread(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     const struct htp_conv2d_hmx_context * hc = (const struct htp_conv2d_hmx_context *) data;
 
     if (hc->unpack_chunk >= 0) {
-        conv2d_hmx_unpack(hc, ith);
+        conv2d_hmx_unpack(hc, th, ith);
     }
     // weights of this rt block are packed before its first chunk
     if (hc->pack_chunk == 0) {
@@ -960,7 +960,7 @@ static void conv2d_hmx_phase_thread(unsigned int nth, unsigned int ith, void * d
         }
     }
     if (hc->pack_chunk >= 0) {
-        conv2d_hmx_pack(hc, ith);
+        conv2d_hmx_pack(hc, th, ith);
     }
 }
 
@@ -978,7 +978,7 @@ static void conv2d_hmx_run_phase(struct htp_conv2d_hmx_context * hc, int32_t pac
     hc->unpack_chunk = unpack;
     hc->pack_tpt     = conv2d_hmx_tpt(hc, pack, nth);
     hc->unpack_tpt   = conv2d_hmx_tpt(hc, unpack, nth);
-    work_queue_run(hc->octx->ctx->work_queue, conv2d_hmx_phase_thread, hc, nth);
+    work_queue_run(hc->octx->ctx->hvx_queue, conv2d_hmx_phase_thread, hc, nth);
 }
 
 static int conv2d_hmx_tiled(struct htp_ops_context * octx, const struct htp_conv2d_kernel_params * kp) {
@@ -997,7 +997,7 @@ static int conv2d_hmx_tiled(struct htp_ops_context * octx, const struct htp_conv
     const size_t out_size = (size_t) kp->rt_block * kp->ct_chunk * tile;
 
     struct htp_conv2d_vtcm_layout layout;
-    htp_conv2d_vtcm_layout_build(&layout, kp->rt_block, w_size, act_size, out_size, kp->thr_size, kp->n_threads);
+    htp_conv2d_vtcm_layout_build(&layout, kp->rt_block, w_size, act_size, out_size, kp->thr_size, kp->n_hvx);
 
     uint8_t * base = (uint8_t *) octx->ctx->vtcm_base;
 
@@ -1015,7 +1015,7 @@ static int conv2d_hmx_tiled(struct htp_ops_context * octx, const struct htp_conv
         .t_end    = (row_base + n_rows) * kp->tiled.n_ow_tiles,
     };
 
-    const uint32_t nth      = kp->n_threads;
+    const uint32_t nth      = kp->n_hvx;
     const int32_t  n_chunks = (int32_t) ((hc.t_end - hc.t_base + kp->ct_chunk - 1) / kp->ct_chunk);
 
     struct htp_conv2d_hmx_task task;
@@ -1023,7 +1023,7 @@ static int conv2d_hmx_tiled(struct htp_ops_context * octx, const struct htp_conv
     for (uint32_t rt0 = 0; rt0 < kp->n_row_tiles; rt0 += kp->rt_block) {
         hc.rt0  = rt0;
         hc.n_rt = MIN(kp->rt_block, kp->n_row_tiles - rt0);
-        work_queue_run(octx->ctx->work_queue, kp->w_kmajor ? conv2d_hmx_wpack_kmajor_thread : conv2d_hmx_wpack_thread, &hc, nth);
+        work_queue_run(octx->ctx->hvx_queue, kp->w_kmajor ? conv2d_hmx_wpack_kmajor_thread : conv2d_hmx_wpack_thread, &hc, nth);
 
         conv2d_hmx_run_phase(&hc, 0, -1, nth);
 
@@ -1110,7 +1110,7 @@ static inline uint32_t conv2d_hmx_win_act_off(uint32_t u) {
     return ((u >> 5) << 11) | ((u & 30) << 6) | ((u & 1) << 1);
 }
 
-static void conv2d_hmx_win_worker(void * data) {
+static void conv2d_hmx_win_worker(const htp_thread_t * th, void * data) {
     const struct htp_conv2d_win_task * j = (const struct htp_conv2d_win_task *) data;
 
     const uint32_t * taps    = j->taps;
@@ -1327,7 +1327,7 @@ static inline void conv2d_hmx_win_wpack_kmajor_compute(const struct htp_conv2d_w
 }
 
 // Cout-contiguous kernel: DMA w[t][c][32 cout] rows, tile vector p = (row 2p, row 2p+1) interleaved per cout
-static void conv2d_hmx_win_wpack_kmajor(const struct htp_conv2d_win_context * wc, unsigned int nth, unsigned int ith) {
+static void conv2d_hmx_win_wpack_kmajor(const struct htp_conv2d_win_context * wc, const htp_thread_t * th, unsigned int nth, unsigned int ith) {
     struct htp_ops_context *             octx   = wc->octx;
     const struct htp_conv2d_kernel_params * kp     = wc->kp;
     const struct htp_tensor *            kernel = octx->src[0];
@@ -1342,13 +1342,13 @@ static void conv2d_hmx_win_wpack_kmajor(const struct htp_conv2d_win_context * wc
     const bool     is_f16 = kernel->type == HTP_TYPE_F16;
     const uint32_t es     = is_f16 ? sizeof(__fp16) : sizeof(float);
 
-    dma_queue *  dmaq       = octx->ctx->dma[ith];
+    dma_queue *  dmaq       = th->dma_q;
     const size_t w_raw_size = (size_t) TG * 32 * VLEN;
     uint8_t *    raws[2]    = { wc->thr_base + (size_t) ith * kp->thr_size,
                                 wc->thr_base + (size_t) ith * kp->thr_size + w_raw_size };
     const HVX_Vector zero   = Q6_V_vzero();
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     const uint32_t n_tg    = (T + TG - 1) / TG;
     const uint32_t n_units = wc->n_rt * n_cc * n_tg;
@@ -1491,7 +1491,7 @@ static void conv2d_hmx_win_wrows_pack(const struct htp_conv2d_win_context * wc,
     }
 }
 
-static void conv2d_hmx_win_wpack_rows(const struct htp_conv2d_win_context * wc, unsigned int nth, unsigned int ith) {
+static void conv2d_hmx_win_wpack_rows(const struct htp_conv2d_win_context * wc, const htp_thread_t * th, unsigned int nth, unsigned int ith) {
     struct htp_ops_context *             octx   = wc->octx;
     const struct htp_conv2d_kernel_params * kp     = wc->kp;
     const struct htp_tensor *            kernel = octx->src[0];
@@ -1501,7 +1501,7 @@ static void conv2d_hmx_win_wpack_rows(const struct htp_conv2d_win_context * wc,
     const uint32_t kpw = kernel->type == HTP_TYPE_F16 ? 64 : 32; // k per vector
 
     const size_t   wraw    = (size_t) 32 * ((CS * T + kpw - 1) / kpw) * VLEN;
-    dma_queue *    dmaq    = octx->ctx->dma[ith];
+    dma_queue *    dmaq    = th->dma_q;
     uint8_t *      raws[2] = { wc->thr_base + (size_t) ith * kp->thr_size, wc->thr_base + (size_t) ith * kp->thr_size + wraw };
     uint8_t *      xr      = raws[1] + wraw;
     HVX_Vector *   vin     = (HVX_Vector *) (xr + wraw);
@@ -1510,7 +1510,7 @@ static void conv2d_hmx_win_wpack_rows(const struct htp_conv2d_win_context * wc,
     const uint32_t n_cs    = (32 + CS - 1) / CS;
     const uint32_t n_units = wc->n_rt * kp->window.n_cc * n_cs;
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     // one DMA per unit: push(N), push(N + 1), pop(N), pack(N), push(N + 2), ...
     for (uint32_t k = 0, unit = ith; k < 2 && unit < n_units; k++, unit += nth) {
@@ -1531,13 +1531,13 @@ static void conv2d_hmx_win_wpack_rows(const struct htp_conv2d_win_context * wc,
     }
 }
 
-static void conv2d_hmx_win_wpack_thread(unsigned int nth, unsigned int ith, void * data) {
+static void conv2d_hmx_win_wpack_thread(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     const struct htp_conv2d_win_context * wc = (const struct htp_conv2d_win_context *) data;
 
     if (wc->kp->w_kmajor) {
-        conv2d_hmx_win_wpack_kmajor(wc, nth, ith);
+        conv2d_hmx_win_wpack_kmajor(wc, th, nth, ith);
     } else {
-        conv2d_hmx_win_wpack_rows(wc, nth, ith);
+        conv2d_hmx_win_wpack_rows(wc, th, nth, ith);
     }
 }
 
@@ -1983,7 +1983,7 @@ static void conv2d_hmx_win_croutons_fold_cwhn(const struct htp_conv2d_win_contex
     }
 }
 
-static void conv2d_hmx_win_pack(const struct htp_conv2d_win_context * wc, unsigned int nth, unsigned int ith) {
+static void conv2d_hmx_win_pack(const struct htp_conv2d_win_context * wc, const htp_thread_t * th, unsigned int nth, unsigned int ith) {
     struct htp_ops_context *             octx = wc->octx;
     const struct htp_conv2d_kernel_params * kp   = wc->kp;
 
@@ -1994,7 +1994,7 @@ static void conv2d_hmx_win_pack(const struct htp_conv2d_win_context * wc, unsign
     const uint32_t b0  = MIN(ith * bpt, ncb);
     const uint32_t b1  = MIN(b0 + bpt, ncb);
 
-    dma_queue * dmaq      = octx->ctx->dma[ith];
+    dma_queue * dmaq      = th->dma_q;
     uint8_t *   thr       = wc->thr_base + (size_t) ith * kp->thr_size;
     float *     stages[2] = { (float *) thr, (float *) (thr + kp->stage_size) };
     __fp16 *    act       = wc->act[wc->pack_chunk & 1];
@@ -2002,7 +2002,7 @@ static void conv2d_hmx_win_pack(const struct htp_conv2d_win_context * wc, unsign
     uint32_t (*stage_fn)(const struct htp_conv2d_win_context *, dma_queue *, float *, uint32_t, uint32_t) =
         kp->window.fold ? conv2d_hmx_win_stage_fold : conv2d_hmx_win_stage;
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     const uint32_t SB = kp->window.seg_b;
 
@@ -2041,7 +2041,7 @@ static void conv2d_hmx_win_pack(const struct htp_conv2d_win_context * wc, unsign
     }
 }
 
-static void conv2d_hmx_win_unpack(const struct htp_conv2d_win_context * wc, unsigned int nth, unsigned int ith) {
+static void conv2d_hmx_win_unpack(const struct htp_conv2d_win_context * wc, const htp_thread_t * th, unsigned int nth, unsigned int ith) {
     struct htp_ops_context *             octx = wc->octx;
     const struct htp_conv2d_kernel_params * kp   = wc->kp;
     const struct htp_tensor *            dst  = octx->dst;
@@ -2060,12 +2060,12 @@ static void conv2d_hmx_win_unpack(const struct htp_conv2d_win_context * wc, unsi
     const uint32_t i0  = MIN(ith * tpt, nt);
     const uint32_t i1  = MIN(i0 + tpt, nt);
 
-    dma_queue *    dmaq      = octx->ctx->dma[ith];
+    dma_queue *    dmaq      = th->dma_q;
     uint8_t *      thr       = wc->thr_base + (size_t) ith * kp->thr_size + 2 * (size_t) kp->stage_size;
     float *        obuf      = (float *) thr;
     __fp16 *       out       = wc->out[wc->unpack_chunk & 1];
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     // nd[bi] = DMAs pushed from obuf bi, popped before it is reused HTP_CONV2D_WIN_OBUF steps later
     uint32_t nd[HTP_CONV2D_WIN_OBUF] = { 0 };
@@ -2154,11 +2154,11 @@ static void conv2d_hmx_win_unpack(const struct htp_conv2d_win_context * wc, unsi
 }
 
 // per-thread callback for one pipeline phase
-static void conv2d_hmx_win_phase_thread(unsigned int nth, unsigned int ith, void * data) {
+static void conv2d_hmx_win_phase_thread(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     const struct htp_conv2d_win_context * wc = (const struct htp_conv2d_win_context *) data;
 
     if (wc->unpack_chunk >= 0) {
-        conv2d_hmx_win_unpack(wc, nth, ith);
+        conv2d_hmx_win_unpack(wc, th, nth, ith);
     }
     // weights of this rt block are packed before its first chunk
     if (wc->pack_chunk == 0) {
@@ -2170,12 +2170,12 @@ static void conv2d_hmx_win_phase_thread(unsigned int nth, unsigned int ith, void
         }
     }
     if (wc->pack_chunk >= 0) {
-        conv2d_hmx_win_pack(wc, nth, ith);
+        conv2d_hmx_win_pack(wc, th, nth, ith);
     }
 }
 
 // zero out 2 act buffer regions in VTCM before packing
-static void conv2d_hmx_win_zero_thread(unsigned int nth, unsigned int ith, void * data) {
+static void conv2d_hmx_win_zero_thread(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     const struct htp_conv2d_win_context * wc = (const struct htp_conv2d_win_context *) data;
 
     const size_t n   = 2 * wc->layout.act_size / sizeof(float);
@@ -2183,7 +2183,7 @@ static void conv2d_hmx_win_zero_thread(unsigned int nth, unsigned int ith, void
     const size_t i0  = MIN((size_t) ith * per, n);
     const size_t i1  = MIN(i0 + per, n);
 
-    struct htp_thread_trace * tr = &wc->octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
     const uint16_t block = (uint16_t) (i0 / VLEN_FP32);
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, block);
     if (i1 > i0) {
@@ -2195,7 +2195,7 @@ static void conv2d_hmx_win_zero_thread(unsigned int nth, unsigned int ith, void
 static void conv2d_hmx_win_run_phase(struct htp_conv2d_win_context * wc, int32_t pack, int32_t unpack) {
     wc->pack_chunk   = pack;
     wc->unpack_chunk = unpack;
-    work_queue_run(wc->octx->ctx->work_queue, conv2d_hmx_win_phase_thread, wc, wc->kp->n_threads);
+    work_queue_run(wc->octx->ctx->hvx_queue, conv2d_hmx_win_phase_thread, wc, wc->kp->n_hvx);
 }
 
 
@@ -2252,7 +2252,7 @@ static int conv2d_hmx_win(struct htp_ops_context * octx, const struct htp_conv2d
     const size_t out_size = (size_t) kp->rt_block * kp->ct_chunk * tile;
 
     struct htp_conv2d_vtcm_layout layout;
-    htp_conv2d_vtcm_layout_build(&layout, kp->rt_block, w_size, act_size, out_size, kp->thr_size, kp->n_threads);
+    htp_conv2d_vtcm_layout_build(&layout, kp->rt_block, w_size, act_size, out_size, kp->thr_size, kp->n_hvx);
 
     uint8_t * base = (uint8_t *) octx->ctx->vtcm_base;
 
@@ -2275,7 +2275,7 @@ static int conv2d_hmx_win(struct htp_ops_context * octx, const struct htp_conv2d
     const uint32_t n_ch = octx->src[1]->ne[2] * (kp->window.fold ? KW : 1);
     if (!kp->cwhn && n_ch <= 2 * CONV2D_WIN_SCAT_MAX) {
         wctx.n_scat = (n_ch + 1) / 2;
-        work_queue_run(octx->ctx->work_queue, conv2d_hmx_win_zero_thread, &wctx, kp->n_threads);
+        work_queue_run(octx->ctx->hvx_queue, conv2d_hmx_win_zero_thread, &wctx, kp->n_hvx);
     }
 
     const int32_t  n_chunks = (int32_t) ((wctx.t_end - wctx.t_base + kp->ct_chunk - 1) / kp->ct_chunk);
@@ -2307,7 +2307,7 @@ static int conv2d_hmx_win(struct htp_ops_context * octx, const struct htp_conv2d
     for (uint32_t rt0 = 0; rt0 < kp->n_row_tiles; rt0 += kp->rt_block) {
         wctx.rt0  = rt0;
         wctx.n_rt = MIN(kp->rt_block, kp->n_row_tiles - rt0);
-        work_queue_run(octx->ctx->work_queue, conv2d_hmx_win_wpack_thread, &wctx, kp->n_threads);
+        work_queue_run(octx->ctx->hvx_queue, conv2d_hmx_win_wpack_thread, &wctx, kp->n_hvx);
 
         conv2d_hmx_win_run_phase(&wctx, 0, -1);
 
@@ -3103,10 +3103,10 @@ struct conv2d_dw_dma_context {
     conv2d_dw_compute_fn_t                     compute_fn;
     uint32_t                                   task_start;
     uint32_t                                   task_count;
-    uint32_t                                   tasks_per_thread;
+    uint32_t                                   tasks_per_hvx;
 };
 
-static void conv2d_dw_dma_thread(unsigned int nth, unsigned int ith, void * data) {
+static void conv2d_dw_dma_thread(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     const struct conv2d_dw_dma_context * ctx  = (const struct conv2d_dw_dma_context *) data;
     const struct htp_ops_context * octx       = ctx->octx;
     const struct htp_conv2d_dw_kernel_params * kp = ctx->kp;
@@ -3115,24 +3115,24 @@ static void conv2d_dw_dma_thread(unsigned int nth, unsigned int ith, void * data
     const struct htp_tensor * dst             = octx->dst;
     conv2d_dw_compute_fn_t compute_fn         = ctx->compute_fn;
 
-    const uint32_t first = ctx->task_start + ith * ctx->tasks_per_thread;
-    const uint32_t last  = MIN(first + ctx->tasks_per_thread, ctx->task_start + ctx->task_count);
+    const uint32_t first = ctx->task_start + ith * ctx->tasks_per_hvx;
+    const uint32_t last  = MIN(first + ctx->tasks_per_hvx, ctx->task_start + ctx->task_count);
     if (first >= last) {
         return;
     }
 
     (void) nth;
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    struct htp_thread_trace * tr = th->trace;
+    dma_queue * dma_q = th->dma_q;
 
     const uint32_t src_spad_half  = (uint32_t) L->src_spad_half_size;
     const uint32_t dst_spad_half  = (uint32_t) L->dst_spad_half_size;
     const uint32_t pad_top_stride = kp->pad_top_stride;
 
     uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
-    uint8_t * src_spad  = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L->off_src) + ith * L->src_bytes_per_thread;
-    uint8_t * dst_spad  = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L->off_dst) + ith * L->dst_bytes_per_thread;
+    uint8_t * src_spad  = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L->off_src) + ith * L->src_bytes_per_hvx;
+    uint8_t * dst_spad  = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L->off_dst) + ith * L->dst_bytes_per_hvx;
 
     float * srcb2[2] = { (float *) src_spad, (float *) (src_spad + src_spad_half) };
     float * dstb2[2] = { (float *) dst_spad, (float *) (dst_spad + dst_spad_half) };
@@ -3213,7 +3213,7 @@ int op_conv_2d_dw(struct htp_ops_context * octx) {
     const uint32_t w_size          = src0->size;
 
     struct htp_conv2d_dw_vtcm_layout layout;
-    htp_conv2d_dw_vtcm_layout_build(&layout, max_src_bytes, max_dst_bytes, w_size, kp->n_threads);
+    htp_conv2d_dw_vtcm_layout_build(&layout, max_src_bytes, max_dst_bytes, w_size, kp->n_hvx);
 
     uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
     void *    w_vtcm    = VTCM_LAYOUT_PTR(void, vtcm_base, layout.off_w);
@@ -3240,8 +3240,8 @@ int op_conv_2d_dw(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    const uint32_t nth = MIN(kp->n_threads, task_count);
-    const uint32_t tasks_per_thread = (task_count + nth - 1) / nth;
+    const uint32_t nth = MIN(kp->n_hvx, task_count);
+    const uint32_t tasks_per_hvx = (task_count + nth - 1) / nth;
 
     struct conv2d_dw_dma_context ctx = {
         .octx             = octx,
@@ -3251,9 +3251,9 @@ int op_conv_2d_dw(struct htp_ops_context * octx) {
         .compute_fn       = conv2d_dw_kernels[kp->kernel_type][kp->is_f16 ? 1 : 0],
         .task_start       = task_start,
         .task_count       = task_count,
-        .tasks_per_thread = tasks_per_thread,
+        .tasks_per_hvx    = tasks_per_hvx,
     };
 
-    work_queue_run(octx->ctx->work_queue, conv2d_dw_dma_thread, &ctx, nth);
+    work_queue_run(octx->ctx->hvx_queue, conv2d_dw_dma_thread, &ctx, nth);
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/conv2d-ops.h

@@ -60,7 +60,7 @@ struct htp_conv2d_window_params {
 
 struct htp_conv2d_kernel_params {
     uint8_t  kernel_type;
-    uint8_t  n_threads;
+    uint8_t  n_hvx;
     uint8_t  cwhn;
     uint8_t  w_kmajor;
 
@@ -98,7 +98,7 @@ enum htp_conv2d_dw_kernel_type {
 
 struct htp_conv2d_dw_kernel_params {
     uint8_t  kernel_type;
-    uint8_t  n_threads;
+    uint8_t  n_hvx;
     uint8_t  cwhn;
     uint8_t  is_f16;
 
@@ -140,8 +140,8 @@ struct htp_conv2d_dw_vtcm_layout {
     size_t   off_dst;
     size_t   src_spad_half_size;
     size_t   dst_spad_half_size;
-    size_t   src_bytes_per_thread;
-    size_t   dst_bytes_per_thread;
+    size_t   src_bytes_per_hvx;
+    size_t   dst_bytes_per_hvx;
 };
 
 static inline void htp_conv2d_dw_vtcm_layout_build(
@@ -149,17 +149,17 @@ static inline void htp_conv2d_dw_vtcm_layout_build(
     uint32_t max_src_bytes,
     uint32_t max_dst_bytes,
     uint32_t w_size,
-    uint32_t n_threads
+    uint32_t n_hvx
 ) {
     L->src_spad_half_size   = hex_round_up(max_src_bytes + 256, 128);
     L->dst_spad_half_size   = hex_round_up(max_dst_bytes + 128, 128);
-    L->src_bytes_per_thread = 2 * L->src_spad_half_size;
-    L->dst_bytes_per_thread = 2 * L->dst_spad_half_size;
+    L->src_bytes_per_hvx    = 2 * L->src_spad_half_size;
+    L->dst_bytes_per_hvx    = 2 * L->dst_spad_half_size;
 
     size_t off = 0;
     VTCM_LAYOUT_ALLOC(off, off_w, hex_round_up(w_size, 128));
-    VTCM_LAYOUT_ALLOC(off, off_src, L->src_bytes_per_thread * n_threads);
-    VTCM_LAYOUT_ALLOC(off, off_dst, L->dst_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_src, L->src_bytes_per_hvx * n_hvx);
+    VTCM_LAYOUT_ALLOC(off, off_dst, L->dst_bytes_per_hvx * n_hvx);
     L->total_bytes = off;
 }
 
@@ -181,7 +181,7 @@ static inline void htp_conv2d_vtcm_layout_build(
     size_t   act_size,
     size_t   out_size,
     size_t   thr_size,
-    uint32_t n_threads
+    uint32_t n_hvx
 ) {
     const size_t sc = hex_round_up((size_t) rt_block * HTP_CONV2D_HMX_SCALES_SIZE, 2048);
     size_t off = 0;
@@ -191,7 +191,7 @@ static inline void htp_conv2d_vtcm_layout_build(
     VTCM_LAYOUT_ALLOC(off, off_act[1], act_size);
     VTCM_LAYOUT_ALLOC(off, off_out[0], out_size);
     VTCM_LAYOUT_ALLOC(off, off_out[1], out_size);
-    VTCM_LAYOUT_ALLOC(off, off_thr,    (size_t) n_threads * thr_size);
+    VTCM_LAYOUT_ALLOC(off, off_thr,    (size_t) n_hvx * thr_size);
     L->act_size    = act_size;
     L->total_bytes = off;
 }

ggml/src/ggml-hexagon/htp/cpy-ops.c

@@ -26,11 +26,11 @@ struct htp_copy_context {
 
     uint32_t row_start;
     uint32_t nrows;
-    uint32_t src0_nrows_per_thread;
+    uint32_t src0_nrows_per_hvx;
 
     uint32_t elem_start;
     uint32_t nelem;
-    uint32_t elem_per_thread;
+    uint32_t elem_per_hvx;
 
     struct htp_copy_convert_vtcm_layout layout;
     uint8_t * vtcm_src0;
@@ -38,7 +38,7 @@ struct htp_copy_context {
 
     uint32_t tile_start;
     uint32_t ntiles;
-    uint32_t tiles_per_thread;
+    uint32_t tiles_per_hvx;
 };
 
 #define cpy_preamble                              \
@@ -66,16 +66,17 @@ struct htp_copy_context {
     const uint32_t  nb3 = dst->nb[3];
 
 #define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                                \
-static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void * data) {            \
+static void cpy_thread_##NAME##_reshape(const htp_thread_t * th, unsigned int nth, unsigned int ith,  \
+                                        void * data) {                                                \
     struct htp_copy_context * ct = (struct htp_copy_context *) data;                                  \
     struct htp_ops_context * octx = ct->octx;                                                         \
     cpy_preamble;                                                                                     \
-    const uint32_t th_nelem = ct->elem_per_thread;                                                    \
+    const uint32_t th_nelem = ct->elem_per_hvx;                                                       \
     const uint32_t th_start = ct->elem_start + ith * th_nelem;                                        \
     const uint32_t th_end   = MIN(th_start + th_nelem, ct->elem_start + ct->nelem);                   \
     if (th_start >= th_end) return;                                                                   \
                                                                                                       \
-    dma_queue * dma_q = octx->ctx->dma[ith];                                                          \
+    dma_queue * dma_q = th->dma_q;                                                                    \
                                                                                                       \
     const uint32_t ne01_ne00      = ne01 * ne00;                                                      \
     const uint32_t ne02_ne01_ne00 = ne02 * ne01_ne00;                                                 \
@@ -215,19 +216,19 @@ static inline void cpy_transpose_push_in(dma_queue * q, const struct htp_copy_tr
 }
 
 // Pipelined tile worker: DMA in -> vgather transpose -> DMA out
-static void cpy_thread_transpose(unsigned int nth, unsigned int ith, void * data) {
+static void cpy_thread_transpose(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct htp_copy_context * ct = (struct htp_copy_context *) data;
     struct htp_ops_context * octx = ct->octx;
     const struct htp_copy_transpose_params * tp = &ct->kparams->u.transpose;
 
-    const uint32_t t0 = ct->tile_start + ith * ct->tiles_per_thread;
-    const uint32_t t1 = MIN(t0 + ct->tiles_per_thread, ct->tile_start + ct->ntiles);
+    const uint32_t t0 = ct->tile_start + ith * ct->tiles_per_hvx;
+    const uint32_t t1 = MIN(t0 + ct->tiles_per_hvx, ct->tile_start + ct->ntiles);
     if (t0 >= t1) {
         return;
     }
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    dma_queue * dma_q = th->dma_q;
+    struct htp_thread_trace * tr = th->trace;
 
     const dma_addr_t src = octx->src[0]->data;
     const dma_addr_t dst = octx->dst->data;
@@ -332,18 +333,19 @@ static void cpy_thread_transpose(unsigned int nth, unsigned int ith, void * data
 }
 
 #define DEFINE_CPY_CONVERT_SAMESHAPE(NAME, CONV_FUNC)                                        \
-static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, void * data) { \
+static void cpy_thread_##NAME##_sameshape(const htp_thread_t * th, unsigned int nth,         \
+                                          unsigned int ith, void * data) {                   \
     struct htp_copy_context * ct = (struct htp_copy_context *) data;                         \
     struct htp_ops_context * octx = ct->octx;                                                \
     cpy_preamble;                                                                            \
                                                                                              \
-    const uint32_t dr  = ct->src0_nrows_per_thread;                                          \
+    const uint32_t dr  = ct->src0_nrows_per_hvx;                                             \
     const uint32_t ir0 = ct->row_start + dr * ith;                                           \
     const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);                           \
     if (ir0 >= ir1) return;                                                                  \
                                                                                              \
-    dma_queue * dma_q = octx->ctx->dma[ith];                                                 \
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                   \
+    dma_queue * dma_q = th->dma_q;                                                           \
+    struct htp_thread_trace * tr = th->trace;                                                \
                                                                                              \
     const struct htp_copy_convert_params * cvt = &ct->kparams->u.convert;                    \
     const struct htp_copy_convert_vtcm_layout * L = &ct->layout;                             \
@@ -352,8 +354,8 @@ static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, vo
     const uint32_t src0_row_stride = L->src0_row_stride;                                     \
     const uint32_t dst_row_stride  = L->dst_row_stride;                                      \
     const uint32_t blk_rows        = L->blk_rows;                                            \
-    uint8_t * vtcm_src0_base = ct->vtcm_src0 + ith * L->spad0_size_per_thread;               \
-    uint8_t * vtcm_dst_base  = ct->vtcm_dst  + ith * L->spad1_size_per_thread;               \
+    uint8_t * vtcm_src0_base = ct->vtcm_src0 + ith * L->spad0_size_per_hvx;                  \
+    uint8_t * vtcm_dst_base  = ct->vtcm_dst  + ith * L->spad1_size_per_hvx;                  \
     const uint32_t src0_row_size = ne00 * ct->kparams->src0_type_size;                       \
     const uint32_t dst_row_size  = ne00 * ct->kparams->dst_type_size;                        \
                                                                                              \
@@ -510,7 +512,7 @@ static int cpy_sameshape_convert(struct htp_ops_context * octx, const struct htp
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
 
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+    if (!htp_ops_context_set_n_hvx(octx, kparams->n_hvx)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
@@ -532,17 +534,17 @@ static int cpy_sameshape_convert(struct htp_ops_context * octx, const struct htp
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_hvx = octx->n_hvx;
     const struct htp_copy_convert_params * cvt = &kparams->u.convert;
 
     struct htp_copy_context ct;
-    htp_copy_convert_vtcm_layout_build(&ct.layout, src0->ne[0], kparams->src0_type_size, kparams->dst_type_size, cvt->blk_rows, n_threads);
+    htp_copy_convert_vtcm_layout_build(&ct.layout, src0->ne[0], kparams->src0_type_size, kparams->dst_type_size, cvt->blk_rows, n_hvx);
 
     ct.octx                  = octx;
     ct.kparams               = kparams;
     ct.row_start             = row_start;
     ct.nrows                 = nrows;
-    ct.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    ct.src0_nrows_per_hvx    = fastdiv(nrows + n_hvx - 1, &octx->n_hvx_div);
 
     uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
     ct.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, ct.layout.off_src0);
@@ -561,7 +563,7 @@ static int cpy_sameshape_convert(struct htp_ops_context * octx, const struct htp
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    work_queue_run(octx->ctx->work_queue, copy_fun, &ct, n_threads);
+    work_queue_run(octx->ctx->hvx_queue, copy_fun, &ct, n_hvx);
     return HTP_STATUS_OK;
 }
 
@@ -569,7 +571,7 @@ static int cpy_reshape(struct htp_ops_context * octx, const struct htp_copy_kern
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
 
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+    if (!htp_ops_context_set_n_hvx(octx, kparams->n_hvx)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
@@ -587,14 +589,14 @@ static int cpy_reshape(struct htp_ops_context * octx, const struct htp_copy_kern
         return HTP_STATUS_OK;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_hvx = octx->n_hvx;
 
     struct htp_copy_context ct;
     ct.octx            = octx;
     ct.kparams         = kparams;
     ct.elem_start      = elem_start;
     ct.nelem           = nelem;
-    ct.elem_per_thread = fastdiv(nelem + n_threads - 1, &octx->n_threads_div);
+    ct.elem_per_hvx    = fastdiv(nelem + n_hvx - 1, &octx->n_hvx_div);
 
     work_queue_func_t copy_fun = NULL;
     switch (src0->type) {
@@ -604,12 +606,12 @@ static int cpy_reshape(struct htp_ops_context * octx, const struct htp_copy_kern
         default: return HTP_STATUS_NO_SUPPORT;
     }
 
-    work_queue_run(octx->ctx->work_queue, copy_fun, &ct, n_threads);
+    work_queue_run(octx->ctx->hvx_queue, copy_fun, &ct, n_hvx);
     return HTP_STATUS_OK;
 }
 
 static int cpy_transpose(struct htp_ops_context * octx, const struct htp_copy_kernel_params * kparams) {
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+    if (!htp_ops_context_set_n_hvx(octx, kparams->n_hvx)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
@@ -633,20 +635,20 @@ static int cpy_transpose(struct htp_ops_context * octx, const struct htp_copy_ke
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_hvx = octx->n_hvx;
 
     struct htp_copy_context ct;
-    ct.octx             = octx;
-    ct.kparams          = kparams;
-    ct.tile_start       = tile_start;
-    ct.ntiles           = ntiles;
-    ct.tiles_per_thread = fastdiv(ntiles + n_threads - 1, &octx->n_threads_div);
+    ct.octx          = octx;
+    ct.kparams       = kparams;
+    ct.tile_start    = tile_start;
+    ct.ntiles        = ntiles;
+    ct.tiles_per_hvx = fastdiv(ntiles + n_hvx - 1, &octx->n_hvx_div);
 
     uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
     ct.vtcm_src0 = vtcm_base;
-    ct.vtcm_dst  = vtcm_base + (size_t) n_threads * 2 * tp->in_buf_size;
+    ct.vtcm_dst  = vtcm_base + (size_t) n_hvx * 2 * tp->in_buf_size;
 
-    work_queue_run(octx->ctx->work_queue, cpy_thread_transpose, &ct, n_threads);
+    work_queue_run(octx->ctx->hvx_queue, cpy_thread_transpose, &ct, n_hvx);
     return HTP_STATUS_OK;
 }
 
@@ -674,7 +676,7 @@ int op_cpy(struct htp_ops_context * octx) {
             status = cpy_transpose(octx, kparams);
             break;
         case HTP_COPY_KERNEL_GATHER_ROWS:
-            status = htp_gather_rows(octx, octx->src[0], NULL, &kparams->u.gather, kparams->n_threads);
+            status = htp_gather_rows(octx, octx->src[0], NULL, &kparams->u.gather, kparams->n_hvx);
             break;
         default:
             status = HTP_STATUS_NO_SUPPORT;

ggml/src/ggml-hexagon/htp/cpy-ops.h

@@ -60,7 +60,7 @@ struct htp_copy_kernel_params {
     uint8_t  kernel_type;
     uint8_t  src0_type_size;
     uint8_t  dst_type_size;
-    uint8_t  n_threads;
+    uint8_t  n_hvx;
 
     uint32_t total_elems;
     uint32_t total_rows;
@@ -88,8 +88,8 @@ struct htp_copy_convert_vtcm_layout {
     uint32_t blk_rows;
     uint32_t src0_buf_size;
     uint32_t dst_buf_size;
-    uint32_t spad0_size_per_thread;
-    uint32_t spad1_size_per_thread;
+    uint32_t spad0_size_per_hvx;
+    uint32_t spad1_size_per_hvx;
 };
 
 static inline void htp_copy_convert_vtcm_layout_build(
@@ -98,7 +98,7 @@ static inline void htp_copy_convert_vtcm_layout_build(
     uint32_t src_type_size,
     uint32_t dst_type_size,
     uint32_t blk_rows,
-    uint32_t n_threads) {
+    uint32_t n_hvx) {
 
     const uint32_t src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
     const uint32_t dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);
@@ -108,12 +108,12 @@ static inline void htp_copy_convert_vtcm_layout_build(
     L->blk_rows              = blk_rows;
     L->src0_buf_size         = blk_rows * src0_row_stride;
     L->dst_buf_size          = blk_rows * dst_row_stride;
-    L->spad0_size_per_thread = 2 * L->src0_buf_size;
-    L->spad1_size_per_thread = 2 * L->dst_buf_size;
+    L->spad0_size_per_hvx    = 2 * L->src0_buf_size;
+    L->spad1_size_per_hvx    = 2 * L->dst_buf_size;
 
     size_t off = 0;
-    VTCM_LAYOUT_ALLOC(off, off_src0, (size_t) n_threads * L->spad0_size_per_thread);
-    VTCM_LAYOUT_ALLOC(off, off_dst,  (size_t) n_threads * L->spad1_size_per_thread);
+    VTCM_LAYOUT_ALLOC(off, off_src0, (size_t) n_hvx * L->spad0_size_per_hvx);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  (size_t) n_hvx * L->spad1_size_per_hvx);
     L->total_bytes = off;
 }
 
@@ -124,7 +124,7 @@ static inline void htp_copy_convert_solve_layout(
     uint32_t ne01,
     uint32_t src_type_size,
     uint32_t dst_type_size,
-    uint32_t n_threads) {
+    uint32_t n_hvx) {
 
     const uint32_t src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
     const uint32_t dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);
@@ -138,7 +138,7 @@ static inline void htp_copy_convert_solve_layout(
         blk_rows = 1;
     }
 
-    htp_copy_convert_vtcm_layout_build(layout, ne00, src_type_size, dst_type_size, blk_rows, n_threads);
+    htp_copy_convert_vtcm_layout_build(layout, ne00, src_type_size, dst_type_size, blk_rows, n_hvx);
 }
 
 #define HTP_COPY_TRANSPOSE_BLK_BYTES (64 * 1024)

ggml/src/ggml-hexagon/htp/cumsum-ops.c

@@ -8,7 +8,7 @@
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
 #include "hex-common.h"
-#include "hex-profile.h"
+#include "htp-trace.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
@@ -48,7 +48,7 @@ struct htp_cumsum_context {
     size_t          dst_row_size;
     size_t          src_row_size_aligned;
     size_t          dst_row_size_aligned;
-    uint32_t        rows_per_thread;
+    uint32_t        rows_per_hvx;
     uint32_t        total_rows;
     uint32_t        row_start;
 };
@@ -57,7 +57,7 @@ struct htp_cumsum_context {
     struct htp_cumsum_context * cctx = (struct htp_cumsum_context *) data; \
     struct htp_ops_context *    octx = cctx->octx;                         \
     htp_cumsum_tensors_preamble;                                           \
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
 
 // ---------------------------------------------------------------------------
 // HVX prefix scan helpers
@@ -118,11 +118,11 @@ static inline void hvx_cumsum_row_f32(const float * restrict src, float * restri
 // Per thread worker: Double-buffered DMA
 // ---------------------------------------------------------------------------
 
-static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * data) {
+static void cumsum_thread_f32_dma(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     htp_cumsum_preamble;
 
-    const uint32_t ir0 = cctx->row_start + cctx->rows_per_thread * ith;
-    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_thread, cctx->row_start + cctx->total_rows);
+    const uint32_t ir0 = cctx->row_start + cctx->rows_per_hvx * ith;
+    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_hvx, cctx->row_start + cctx->total_rows);
 
     if (ir0 >= ir1) {
         return;
@@ -151,7 +151,7 @@ static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * dat
                        src_row_size_aligned, src_row_size, src_row_size, 1);
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     for (uint32_t ir = ir0; ir < ir1; ir++) {
         float * dst_spad_row = (float *) dma_queue_pop(dma_q).src;
@@ -184,16 +184,16 @@ static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * dat
 // Per thread worker: Direct HVX (no DMA)
 // ---------------------------------------------------------------------------
 
-static void cumsum_thread_f32(unsigned int nth, unsigned int ith, void * data) {
+static void cumsum_thread_f32(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     htp_cumsum_preamble;
 
     const uint8_t * src_data = (const uint8_t *) src0->data;
     uint8_t *       dst_data = (uint8_t *) dst->data;
 
-    const uint32_t ir0 = cctx->row_start + cctx->rows_per_thread * ith;
-    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_thread, cctx->row_start + cctx->total_rows);
+    const uint32_t ir0 = cctx->row_start + cctx->rows_per_hvx * ith;
+    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_hvx, cctx->row_start + cctx->total_rows);
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
 
     for (uint32_t ir = ir0; ir < ir1; ir++) {
@@ -234,11 +234,11 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+    if (!htp_ops_context_set_n_hvx(octx, MIN(kparams->n_hvx, nrows))) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_hvx = octx->n_hvx;
 
     if (!kparams->use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
         return HTP_STATUS_NO_SUPPORT;
@@ -246,12 +246,12 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
 
     struct htp_cumsum_vtcm_layout layout = { 0 };
     if (kparams->use_dma) {
-        htp_cumsum_vtcm_layout_build(&layout, src_row_size, dst_row_size, n_threads);
+        htp_cumsum_vtcm_layout_build(&layout, src_row_size, dst_row_size, n_hvx);
 
-        octx->src0_spad.size_per_thread = layout.src0_bytes_per_thread;
-        octx->dst_spad.size_per_thread  = layout.dst_bytes_per_thread;
-        octx->src0_spad.size = n_threads * layout.src0_bytes_per_thread;
-        octx->dst_spad.size  = n_threads * layout.dst_bytes_per_thread;
+        octx->src0_spad.size_per_hvx = layout.src0_bytes_per_hvx;
+        octx->dst_spad.size_per_hvx  = layout.dst_bytes_per_hvx;
+        octx->src0_spad.size = n_hvx * layout.src0_bytes_per_hvx;
+        octx->dst_spad.size  = n_hvx * layout.dst_bytes_per_hvx;
         octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_src0);
         octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_dst);
         octx->src0_spad.src  = NULL;
@@ -264,15 +264,15 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
         .dst_row_size         = dst_row_size,
         .src_row_size_aligned = layout.src_row_size_aligned,
         .dst_row_size_aligned = layout.dst_row_size_aligned,
-        .rows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+        .rows_per_hvx         = fastdiv(nrows + n_hvx - 1, &octx->n_hvx_div),
         .total_rows           = nrows,
         .row_start            = row_start,
     };
 
     if (kparams->use_dma) {
-        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32_dma, &cctx, n_threads);
+        work_queue_run(octx->ctx->hvx_queue, cumsum_thread_f32_dma, &cctx, n_hvx);
     } else {
-        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32, &cctx, n_threads);
+        work_queue_run(octx->ctx->hvx_queue, cumsum_thread_f32, &cctx, n_hvx);
     }
 
     return HTP_STATUS_OK;

ggml/src/ggml-hexagon/htp/cumsum-ops.h

@@ -11,30 +11,30 @@ struct htp_cumsum_vtcm_layout {
     size_t off_dst;
     size_t src_row_size_aligned;
     size_t dst_row_size_aligned;
-    size_t src0_bytes_per_thread;
-    size_t dst_bytes_per_thread;
+    size_t src0_bytes_per_hvx;
+    size_t dst_bytes_per_hvx;
 };
 
 static inline void htp_cumsum_vtcm_layout_build(
     struct htp_cumsum_vtcm_layout * L,
     size_t                          src_row_size,
     size_t                          dst_row_size,
-    uint32_t                        n_threads
+    uint32_t                        n_hvx
 ) {
     L->src_row_size_aligned  = hex_round_up(src_row_size, 128);
     L->dst_row_size_aligned  = hex_round_up(dst_row_size, 128);
-    L->src0_bytes_per_thread = 2 * L->src_row_size_aligned;
-    L->dst_bytes_per_thread  = 2 * L->dst_row_size_aligned;
+    L->src0_bytes_per_hvx    = 2 * L->src_row_size_aligned;
+    L->dst_bytes_per_hvx     = 2 * L->dst_row_size_aligned;
 
     size_t off = 0;
-    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
-    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_hvx * n_hvx);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_hvx * n_hvx);
 
     L->total_bytes = off;
 }
 
 struct htp_cumsum_kernel_params {
-    uint32_t n_threads;
+    uint32_t n_hvx;
     uint32_t total_rows;
     uint32_t use_dma;
     uint32_t vtcm_size;

ggml/src/ggml-hexagon/htp/diag-ops.c

@@ -6,7 +6,7 @@
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
 #include "hex-common.h"
-#include "hex-profile.h"
+#include "htp-trace.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
@@ -40,7 +40,7 @@ struct htp_diag_context {
     size_t          dst_row_size;
     size_t          src_batch_size_aligned;
     size_t          dst_row_size_aligned;
-    uint32_t        batches_per_thread;
+    uint32_t        batches_per_hvx;
     uint32_t        total_batches;
     uint32_t        batch_start;
 };
@@ -60,12 +60,12 @@ static inline void hvx_diag_row_f32(const float * restrict src, float * restrict
 // Per thread worker: DMA src fetch, compute in VTCM, DMA dst writeback
 // ---------------------------------------------------------------------------
 
-static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data) {
+static void diag_thread_f32_dma(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     htp_diag_preamble;
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
 
-    const uint32_t ib0 = dctx->batch_start + dctx->batches_per_thread * ith;
-    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->batch_start + dctx->total_batches);
+    const uint32_t ib0 = dctx->batch_start + dctx->batches_per_hvx * ith;
+    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_hvx, dctx->batch_start + dctx->total_batches);
 
     if (ib0 >= ib1) {
         return;
@@ -83,7 +83,7 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
     uint8_t * src_spad = octx->src0_spad.data + (ith * src_batch_size_aligned);
     uint8_t * dst_spad = octx->dst_spad.data  + (ith * dst_row_size_aligned);
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     for (uint32_t ib = ib0; ib < ib1; ib++) {
         const uint32_t i3 = fastdiv(ib, &dctx->kparams->div_ne02);
@@ -124,16 +124,16 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
 // Per thread worker: Direct HVX (no DMA)
 // ---------------------------------------------------------------------------
 
-static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
+static void diag_thread_f32(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     htp_diag_preamble;
 
     const uint8_t * src_data = (const uint8_t *) src0->data;
     uint8_t *       dst_data = (uint8_t *) dst->data;
 
-    const uint32_t ib0 = dctx->batch_start + dctx->batches_per_thread * ith;
-    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->batch_start + dctx->total_batches);
+    const uint32_t ib0 = dctx->batch_start + dctx->batches_per_hvx * ith;
+    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_hvx, dctx->batch_start + dctx->total_batches);
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ib0);
 
     for (uint32_t ib = ib0; ib < ib1; ib++) {
@@ -177,11 +177,11 @@ int op_diag_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nbatches))) {
+    if (!htp_ops_context_set_n_hvx(octx, MIN(kparams->n_hvx, nbatches))) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_hvx = octx->n_hvx;
 
     if (!kparams->use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
         return HTP_STATUS_NO_SUPPORT;
@@ -192,12 +192,12 @@ int op_diag_f32(struct htp_ops_context * octx) {
 
     struct htp_diag_vtcm_layout layout = { 0 };
     if (kparams->use_dma) {
-        htp_diag_vtcm_layout_build(&layout, src_batch_size, dst_row_size, n_threads);
+        htp_diag_vtcm_layout_build(&layout, src_batch_size, dst_row_size, n_hvx);
 
-        octx->src0_spad.size_per_thread = layout.src0_bytes_per_thread;
-        octx->dst_spad.size_per_thread  = layout.dst_bytes_per_thread;
-        octx->src0_spad.size = n_threads * layout.src0_bytes_per_thread;
-        octx->dst_spad.size  = n_threads * layout.dst_bytes_per_thread;
+        octx->src0_spad.size_per_hvx = layout.src0_bytes_per_hvx;
+        octx->dst_spad.size_per_hvx  = layout.dst_bytes_per_hvx;
+        octx->src0_spad.size = n_hvx * layout.src0_bytes_per_hvx;
+        octx->dst_spad.size  = n_hvx * layout.dst_bytes_per_hvx;
         octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_src0);
         octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_dst);
         octx->src0_spad.src  = NULL;
@@ -211,15 +211,15 @@ int op_diag_f32(struct htp_ops_context * octx) {
         .dst_row_size           = dst_row_size,
         .src_batch_size_aligned = layout.src_batch_size_aligned,
         .dst_row_size_aligned   = layout.dst_row_size_aligned,
-        .batches_per_thread     = fastdiv(nbatches + n_threads - 1, &octx->n_threads_div),
+        .batches_per_hvx        = fastdiv(nbatches + n_hvx - 1, &octx->n_hvx_div),
         .total_batches          = nbatches,
         .batch_start            = batch_start,
     };
 
     if (kparams->use_dma) {
-        work_queue_run(octx->ctx->work_queue, diag_thread_f32_dma, &dctx, n_threads);
+        work_queue_run(octx->ctx->hvx_queue, diag_thread_f32_dma, &dctx, n_hvx);
     } else {
-        work_queue_run(octx->ctx->work_queue, diag_thread_f32, &dctx, n_threads);
+        work_queue_run(octx->ctx->hvx_queue, diag_thread_f32, &dctx, n_hvx);
     }
 
     return HTP_STATUS_OK;

ggml/src/ggml-hexagon/htp/diag-ops.h

@@ -12,30 +12,30 @@ struct htp_diag_vtcm_layout {
     size_t off_dst;
     size_t src_batch_size_aligned;
     size_t dst_row_size_aligned;
-    size_t src0_bytes_per_thread;
-    size_t dst_bytes_per_thread;
+    size_t src0_bytes_per_hvx;
+    size_t dst_bytes_per_hvx;
 };
 
 static inline void htp_diag_vtcm_layout_build(
     struct htp_diag_vtcm_layout * L,
     size_t                        src_batch_size,
     size_t                        dst_row_size,
-    uint32_t                      n_threads
+    uint32_t                      n_hvx
 ) {
     L->src_batch_size_aligned = hex_round_up(src_batch_size, 128);
     L->dst_row_size_aligned   = hex_round_up(dst_row_size, 128);
-    L->src0_bytes_per_thread  = L->src_batch_size_aligned;
-    L->dst_bytes_per_thread   = L->dst_row_size_aligned;
+    L->src0_bytes_per_hvx     = L->src_batch_size_aligned;
+    L->dst_bytes_per_hvx      = L->dst_row_size_aligned;
 
     size_t off = 0;
-    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
-    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_hvx * n_hvx);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_hvx * n_hvx);
 
     L->total_bytes = off;
 }
 
 struct htp_diag_kernel_params {
-    uint32_t n_threads;
+    uint32_t n_hvx;
     uint32_t total_batches;
     uint32_t use_dma;
     uint32_t can_split;

ggml/src/ggml-hexagon/htp/dma-queue.h

@@ -9,7 +9,7 @@
 #include <string.h>
 #include "hex-utils.h"
 
-#include "hex-profile.h"
+#include "htp-trace.h"
 
 #ifdef __cplusplus
 extern "C" {

ggml/src/ggml-hexagon/htp/dsv4-hc-ops.c

@@ -9,7 +9,7 @@
 
 #include "dma-queue.h"
 #include "hex-common.h"
-#include "hex-profile.h"
+#include "htp-trace.h"
 #include "htp-ctx.h"
 #include "htp-tensor.h"
 #include "hvx-utils.h"
@@ -57,7 +57,7 @@ struct dsv4_hc_pre_context {
     uint32_t token_start;
     uint32_t n_chunks;
     uint32_t total_jobs;
-    uint32_t jobs_per_thread;
+    uint32_t jobs_per_hvx;
     float scale;
     uint8_t * vtcm;
     uint32_t dma_block_chunks;
@@ -178,7 +178,7 @@ static __attribute__((noinline)) void dsv4_hc_pre_gated_block_f32(
     }
 }
 
-static void dsv4_hc_pre_thread_plain_f32(unsigned int nth, unsigned int ith, void * data) {
+static void dsv4_hc_pre_thread_plain_f32(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct dsv4_hc_pre_context * dctx = (struct dsv4_hc_pre_context *) data;
     struct htp_ops_context *     octx = dctx->octx;
 
@@ -187,8 +187,8 @@ static void dsv4_hc_pre_thread_plain_f32(unsigned int nth, unsigned int ith, voi
 
     const uint32_t row_stride  = dctx->dma_row_stride;
     const uint32_t buffer_size = dctx->dma_buffer_size;
-    const uint32_t start       = dctx->jobs_per_thread * ith;
-    const uint32_t end         = hex_smin(start + dctx->jobs_per_thread, dctx->total_jobs);
+    const uint32_t start       = dctx->jobs_per_hvx * ith;
+    const uint32_t end         = hex_smin(start + dctx->jobs_per_hvx, dctx->total_jobs);
     if (start >= end) {
         return;
     }
@@ -198,7 +198,7 @@ static void dsv4_hc_pre_thread_plain_f32(unsigned int nth, unsigned int ith, voi
         dctx->vtcm + 2*buffer_size*ith + buffer_size,
     };
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     uint32_t prefetch_job = start;
     for (uint32_t i = 0; i < 2 && prefetch_job < end; ++i) {
         const struct dsv4_hc_block block = dsv4_hc_pre_block_make(dctx, prefetch_job, end);
@@ -208,7 +208,7 @@ static void dsv4_hc_pre_thread_plain_f32(unsigned int nth, unsigned int ith, voi
         prefetch_job += block.count;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     uint32_t job  = start;
     uint32_t curr = 0;
@@ -244,7 +244,7 @@ static void dsv4_hc_pre_thread_plain_f32(unsigned int nth, unsigned int ith, voi
     (void) nth;
 }
 
-static void dsv4_hc_pre_thread_gated_f32(unsigned int nth, unsigned int ith, void * data) {
+static void dsv4_hc_pre_thread_gated_f32(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct dsv4_hc_pre_context * dctx = (struct dsv4_hc_pre_context *) data;
     struct htp_ops_context *     octx = dctx->octx;
 
@@ -253,8 +253,8 @@ static void dsv4_hc_pre_thread_gated_f32(unsigned int nth, unsigned int ith, voi
 
     const uint32_t row_stride  = dctx->dma_row_stride;
     const uint32_t buffer_size = dctx->dma_buffer_size;
-    const uint32_t start       = dctx->jobs_per_thread * ith;
-    const uint32_t end         = hex_smin(start + dctx->jobs_per_thread, dctx->total_jobs);
+    const uint32_t start       = dctx->jobs_per_hvx * ith;
+    const uint32_t end         = hex_smin(start + dctx->jobs_per_hvx, dctx->total_jobs);
     if (start >= end) {
         return;
     }
@@ -264,7 +264,7 @@ static void dsv4_hc_pre_thread_gated_f32(unsigned int nth, unsigned int ith, voi
         dctx->vtcm + 2*buffer_size*ith + buffer_size,
     };
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     uint32_t prefetch_job = start;
     for (uint32_t i = 0; i < 2 && prefetch_job < end; ++i) {
         const struct dsv4_hc_block block = dsv4_hc_pre_block_make(dctx, prefetch_job, end);
@@ -274,7 +274,7 @@ static void dsv4_hc_pre_thread_gated_f32(unsigned int nth, unsigned int ith, voi
         prefetch_job += block.count;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     uint32_t job  = start;
     uint32_t curr = 0;
@@ -317,22 +317,22 @@ int op_dsv4_hc_pre(struct htp_ops_context * octx) {
     const struct htp_dsv4_hc_pre_kernel_params * kparams =
         (const struct htp_dsv4_hc_pre_kernel_params *) octx->kernel_params;
 
-    const uint32_t n_threads = kparams->n_threads;
-    if (!htp_ops_context_set_n_threads_div(octx, n_threads, kparams->div_n_threads)) {
+    const uint32_t n_hvx = kparams->n_hvx;
+    if (!htp_ops_context_set_n_hvx_div(octx, n_hvx, kparams->div_n_hvx)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
     uint32_t token_start     = 0;
     uint32_t n_tokens        = x->ne[2];
     uint32_t total_jobs      = kparams->total_jobs;
-    uint32_t jobs_per_thread = kparams->jobs_per_thread;
+    uint32_t jobs_per_hvx    = kparams->jobs_per_hvx;
     if (octx->ctx->mdev.count > 1) {
         const bool can_split = htp_tensor_mdev_data_aligned(dst) && ((dst->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0);
         const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(n_tokens, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         token_start     = range.start;
         n_tokens        = range.count;
         total_jobs      = n_tokens * kparams->n_chunks;
-        jobs_per_thread = fastdiv(total_jobs + n_threads - 1, &kparams->div_n_threads);
+        jobs_per_hvx    = fastdiv(total_jobs + n_hvx - 1, &kparams->div_n_hvx);
     }
 
     if (n_tokens == 0) {
@@ -348,7 +348,7 @@ int op_dsv4_hc_pre(struct htp_ops_context * octx) {
         .token_start      = token_start,
         .n_chunks         = kparams->n_chunks,
         .total_jobs       = total_jobs,
-        .jobs_per_thread  = jobs_per_thread,
+        .jobs_per_hvx     = jobs_per_hvx,
         .scale            = scale,
         .vtcm             = (uint8_t *) octx->ctx->vtcm_base,
         .dma_block_chunks = kparams->dma_block_chunks,
@@ -363,11 +363,11 @@ int op_dsv4_hc_pre(struct htp_ops_context * octx) {
         default: return HTP_STATUS_INVAL_PARAMS;
     }
 
-    FARF(HIGH, "dsv4-hc-pre: (%ux%ux%u) -> (%ux%u) : ktype %u tokens %u:%u nth %u\n",
+    FARF(HIGH, "dsv4-hc-pre: (%ux%ux%u) -> (%ux%u) : ktype %u tokens %u:%u nhvx %u\n",
          x->ne[0], x->ne[1], x->ne[2], dst->ne[0], dst->ne[1], kparams->kernel_type,
-         token_start, token_start + n_tokens, n_threads);
+         token_start, token_start + n_tokens, n_hvx);
 
-    work_queue_run(octx->ctx->work_queue, worker, &dctx, n_threads);
+    work_queue_run(octx->ctx->hvx_queue, worker, &dctx, n_hvx);
     return HTP_STATUS_OK;
 }
 
@@ -380,7 +380,7 @@ struct dsv4_hc_post_context {
     const struct htp_dsv4_hc_post_kernel_params * kparams;
     uint32_t token_start;
     uint32_t total_jobs;
-    uint32_t jobs_per_thread;
+    uint32_t jobs_per_hvx;
     uint8_t * vtcm;
 };
 
@@ -481,7 +481,7 @@ static inline void dsv4_hc_post_push_comb(
                    16, comb->nb[1], 16, 4);
 }
 
-static void dsv4_hc_post_thread_identity_f32(unsigned int nth, unsigned int ith, void * data) {
+static void dsv4_hc_post_thread_identity_f32(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct dsv4_hc_post_context * dctx = (struct dsv4_hc_post_context *) data;
     struct htp_ops_context *      octx = dctx->octx;
 
@@ -491,8 +491,8 @@ static void dsv4_hc_post_thread_identity_f32(unsigned int nth, unsigned int ith,
     const uint32_t block_chunks = dctx->kparams->dma_block_chunks;
     const uint32_t row_stride   = dctx->kparams->dma_row_stride;
     const uint32_t buffer_size  = dctx->kparams->dma_buffer_size;
-    const uint32_t start        = dctx->jobs_per_thread * ith;
-    const uint32_t end          = hex_smin(start + dctx->jobs_per_thread, dctx->total_jobs);
+    const uint32_t start        = dctx->jobs_per_hvx * ith;
+    const uint32_t end          = hex_smin(start + dctx->jobs_per_hvx, dctx->total_jobs);
     if (start >= end) {
         return;
     }
@@ -502,7 +502,7 @@ static void dsv4_hc_post_thread_identity_f32(unsigned int nth, unsigned int ith,
         dctx->vtcm + 2*buffer_size*ith + buffer_size,
     };
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     uint32_t prefetch_job = start;
     for (uint32_t i = 0; i < 2 && prefetch_job < end; ++i) {
         const struct dsv4_hc_block block = dsv4_hc_post_block_make(dctx, prefetch_job, end, block_chunks);
@@ -512,7 +512,7 @@ static void dsv4_hc_post_thread_identity_f32(unsigned int nth, unsigned int ith,
         prefetch_job += block.count;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     uint32_t job  = start;
     uint32_t curr = 0;
@@ -549,7 +549,7 @@ static void dsv4_hc_post_thread_identity_f32(unsigned int nth, unsigned int ith,
     (void) nth;
 }
 
-static void dsv4_hc_post_thread_comb_f32(unsigned int nth, unsigned int ith, void * data) {
+static void dsv4_hc_post_thread_comb_f32(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct dsv4_hc_post_context * dctx = (struct dsv4_hc_post_context *) data;
     struct htp_ops_context *      octx = dctx->octx;
 
@@ -559,8 +559,8 @@ static void dsv4_hc_post_thread_comb_f32(unsigned int nth, unsigned int ith, voi
     const uint32_t block_chunks = dctx->kparams->dma_block_chunks;
     const uint32_t row_stride   = dctx->kparams->dma_row_stride;
     const uint32_t buffer_size  = dctx->kparams->dma_buffer_size;
-    const uint32_t start        = dctx->jobs_per_thread * ith;
-    const uint32_t end          = hex_smin(start + dctx->jobs_per_thread, dctx->total_jobs);
+    const uint32_t start        = dctx->jobs_per_hvx * ith;
+    const uint32_t end          = hex_smin(start + dctx->jobs_per_hvx, dctx->total_jobs);
     if (start >= end) {
         return;
     }
@@ -570,7 +570,7 @@ static void dsv4_hc_post_thread_comb_f32(unsigned int nth, unsigned int ith, voi
         dctx->vtcm + 2*buffer_size*ith + buffer_size,
     };
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     uint32_t prefetch_job = start;
     for (uint32_t i = 0; i < 2 && prefetch_job < end; ++i) {
         const struct dsv4_hc_block block = dsv4_hc_post_block_make(dctx, prefetch_job, end, block_chunks);
@@ -580,7 +580,7 @@ static void dsv4_hc_post_thread_comb_f32(unsigned int nth, unsigned int ith, voi
         prefetch_job += block.count;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     uint32_t job  = start;
     uint32_t curr = 0;
@@ -627,22 +627,22 @@ int op_dsv4_hc_post(struct htp_ops_context * octx) {
     const struct htp_dsv4_hc_post_kernel_params * kparams =
         (const struct htp_dsv4_hc_post_kernel_params *) octx->kernel_params;
 
-    const uint32_t n_threads = kparams->n_threads;
-    if (!htp_ops_context_set_n_threads_div(octx, n_threads, kparams->div_n_threads)) {
+    const uint32_t n_hvx = kparams->n_hvx;
+    if (!htp_ops_context_set_n_hvx_div(octx, n_hvx, kparams->div_n_hvx)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
     uint32_t token_start     = 0;
     uint32_t n_tokens        = x->ne[1];
     uint32_t total_jobs      = kparams->total_jobs;
-    uint32_t jobs_per_thread = kparams->jobs_per_thread;
+    uint32_t jobs_per_hvx    = kparams->jobs_per_hvx;
     if (octx->ctx->mdev.count > 1) {
         const bool can_split = htp_tensor_mdev_data_aligned(dst) && ((dst->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0);
         const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(n_tokens, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         token_start     = range.start;
         n_tokens        = range.count;
         total_jobs      = n_tokens * kparams->n_chunks;
-        jobs_per_thread = fastdiv(total_jobs + n_threads - 1, &kparams->div_n_threads);
+        jobs_per_hvx    = fastdiv(total_jobs + n_hvx - 1, &kparams->div_n_hvx);
     }
 
     if (n_tokens == 0) {
@@ -654,7 +654,7 @@ int op_dsv4_hc_post(struct htp_ops_context * octx) {
         .kparams         = kparams,
         .token_start     = token_start,
         .total_jobs      = total_jobs,
-        .jobs_per_thread = jobs_per_thread,
+        .jobs_per_hvx    = jobs_per_hvx,
         .vtcm            = (uint8_t *) octx->ctx->vtcm_base,
     };
 
@@ -665,11 +665,11 @@ int op_dsv4_hc_post(struct htp_ops_context * octx) {
         default: return HTP_STATUS_INVAL_PARAMS;
     }
 
-    FARF(HIGH, "dsv4-hc-post: (%ux%u) (%ux%ux%u) -> (%ux%ux%u) : ktype %u tokens %u:%u nth %u\n",
+    FARF(HIGH, "dsv4-hc-post: (%ux%u) (%ux%ux%u) -> (%ux%ux%u) : ktype %u tokens %u:%u nhvx %u\n",
          x->ne[0], x->ne[1], residual->ne[0], residual->ne[1], residual->ne[2],
          dst->ne[0], dst->ne[1], dst->ne[2], kparams->kernel_type,
-         token_start, token_start + n_tokens, n_threads);
+         token_start, token_start + n_tokens, n_hvx);
 
-    work_queue_run(octx->ctx->work_queue, worker, &dctx, n_threads);
+    work_queue_run(octx->ctx->hvx_queue, worker, &dctx, n_hvx);
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/dsv4-hc-ops.h

@@ -18,10 +18,10 @@ enum htp_dsv4_hc_pre_kernel_type {
 
 struct htp_dsv4_hc_pre_kernel_params {
     uint32_t kernel_type;
-    uint32_t n_threads;
+    uint32_t n_hvx;
     uint32_t n_chunks;
     uint32_t total_jobs;
-    uint32_t jobs_per_thread;
+    uint32_t jobs_per_hvx;
 
     uint32_t dma_block_chunks;
     uint32_t dma_row_stride;
@@ -29,7 +29,7 @@ struct htp_dsv4_hc_pre_kernel_params {
     uint32_t vtcm_size;
 
     struct fastdiv_values div_n_chunks;
-    struct fastdiv_values div_n_threads;
+    struct fastdiv_values div_n_hvx;
 };
 
 #if defined(__cplusplus)
@@ -48,10 +48,10 @@ enum htp_dsv4_hc_post_kernel_type {
 
 struct htp_dsv4_hc_post_kernel_params {
     uint32_t kernel_type;
-    uint32_t n_threads;
+    uint32_t n_hvx;
     uint32_t n_chunks;
     uint32_t total_jobs;
-    uint32_t jobs_per_thread;
+    uint32_t jobs_per_hvx;
     uint32_t has_comb;
 
     uint32_t dma_block_chunks;
@@ -60,7 +60,7 @@ struct htp_dsv4_hc_post_kernel_params {
     uint32_t vtcm_size;
 
     struct fastdiv_values div_n_chunks;
-    struct fastdiv_values div_n_threads;
+    struct fastdiv_values div_n_hvx;
 };
 
 #if defined(__cplusplus)

ggml/src/ggml-hexagon/htp/fill-ops.c

@@ -6,7 +6,7 @@
 #include <string.h>
 
 #include "hex-common.h"
-#include "hex-profile.h"
+#include "htp-trace.h"
 
 #include "hvx-copy.h"
 #include "hvx-utils.h"
@@ -36,27 +36,27 @@
 struct htp_fill_context {
     struct htp_ops_context * octx;
     const struct htp_fill_kernel_params * kparams;
-    uint32_t nrows_per_thread;
+    uint32_t nrows_per_hvx;
     uint32_t total_rows;  // ne1 * ne2 * ne3
     uint32_t row_start;
     HVX_Vector splat_vec;
     uint32_t   elem_size;
 };
 
-static void fill_thread_contiguous(unsigned int nth, unsigned int ith, void * data) {
+static void fill_thread_contiguous(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     const struct htp_fill_context * fctx = (const struct htp_fill_context *) data;
     struct htp_ops_context        * octx = fctx->octx;
     fill_preamble;
 
     // Parallelise over the flat row index spanning ne1*ne2*ne3
-    const uint32_t ir0 = fctx->row_start + fctx->nrows_per_thread * ith;
-    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_thread, fctx->row_start + fctx->total_rows);
+    const uint32_t ir0 = fctx->row_start + fctx->nrows_per_hvx * ith;
+    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_hvx, fctx->row_start + fctx->total_rows);
 
     if (ir0 >= ir1) {
         return;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
 
     const uint32_t elem_start = ir0 * ne0;
@@ -68,21 +68,21 @@ static void fill_thread_contiguous(unsigned int nth, unsigned int ith, void * da
     FARF(HIGH, "fill-contig %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
 }
 
-static void fill_thread_strided(unsigned int nth, unsigned int ith, void * data) {
+static void fill_thread_strided(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     const struct htp_fill_context * fctx = (const struct htp_fill_context *) data;
     struct htp_ops_context        * octx = fctx->octx;
     const struct htp_fill_kernel_params * kparams = fctx->kparams;
     fill_preamble;
 
     // Parallelise over the flat row index spanning ne1*ne2*ne3
-    const uint32_t ir0 = fctx->row_start + fctx->nrows_per_thread * ith;
-    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_thread, fctx->row_start + fctx->total_rows);
+    const uint32_t ir0 = fctx->row_start + fctx->nrows_per_hvx * ith;
+    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_hvx, fctx->row_start + fctx->total_rows);
 
     if (ir0 >= ir1) {
         return;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
 
     for (uint32_t ir = ir0; ir < ir1; ++ir) {
@@ -120,11 +120,11 @@ int op_fill(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+    if (!htp_ops_context_set_n_hvx(octx, MIN(kparams->n_hvx, nrows))) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_hvx = octx->n_hvx;
 
     FARF(HIGH, "fill: (%ux%ux%ux%u) type=%u ktype=%u\n",
          dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], dst->type, kparams->kernel_type);
@@ -135,7 +135,7 @@ int op_fill(struct htp_ops_context * octx) {
     struct htp_fill_context fctx = {
         .octx             = octx,
         .kparams          = kparams,
-        .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+        .nrows_per_hvx    = fastdiv(nrows + n_hvx - 1, &octx->n_hvx_div),
         .total_rows       = nrows,
         .row_start        = row_start,
         .elem_size        = kparams->elem_size,
@@ -156,7 +156,7 @@ int op_fill(struct htp_ops_context * octx) {
         ? fill_thread_contiguous
         : fill_thread_strided;
 
-    work_queue_run(octx->ctx->work_queue, worker, &fctx, n_threads);
+    work_queue_run(octx->ctx->hvx_queue, worker, &fctx, n_hvx);
 
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/fill-ops.h

@@ -11,7 +11,7 @@ enum htp_fill_kernel_type {
 
 struct htp_fill_kernel_params {
     uint8_t  kernel_type;
-    uint32_t n_threads;
+    uint32_t n_hvx;
     uint32_t total_rows;
     uint32_t elem_size;
     uint32_t rows_per_chunk;

ggml/src/ggml-hexagon/htp/flash-attn-ops.c

@@ -14,7 +14,7 @@
 
 #include "dma-queue.h"
 #include "hex-fastdiv.h"
-#include "hex-profile.h"
+#include "htp-trace.h"
 #include "hmx-queue.h"
 #include "hmx-utils.h"
 #include "hvx-utils.h"
@@ -44,7 +44,7 @@
 struct hvx_fa_ddr_spad {
     __fp16   slopes[512];
     uint16_t active_blocks[FA_MAX_KV_BLOCKS];
-    uint16_t thread_active_blocks[HTP_MAX_NTHREADS][FA_MAX_KV_BLOCKS] __attribute__((aligned(128)));
+    uint16_t thread_active_blocks[HTP_MAX_NHVX][FA_MAX_KV_BLOCKS] __attribute__((aligned(128)));
 };
 
 struct hmx_fa_ddr_spad {
@@ -95,7 +95,7 @@ struct htp_fa_context {
     size_t size_m_block;
 
     uint32_t qrows;
-    uint32_t qrows_per_thread;
+    uint32_t qrows_per_hvx;
     uint32_t qrow_start;
 
     bool is_q_fp32;
@@ -111,7 +111,7 @@ struct htp_fa_context {
     float *    spad_sinks;
 
     struct hvx_fa_ddr_spad * ddr;
-    dma_cache_dm *           m_cache; // [n_threads]
+    dma_cache_dm *           m_cache; // [n_hvx]
 
     const struct htp_tensor * k;
     const struct htp_tensor * v;
@@ -122,8 +122,8 @@ struct hmx_fa_context {
     const struct htp_tensor *      sinks;  // attention sinks (src[4]), NULL if absent
     const struct htp_tensor *      k;
     const struct htp_tensor *      v;
-    bool         pipeline;  // true when n_kv_blocks >= FA_MIN_KV_BLOCKS && n_threads >= 2
-    uint32_t     n_threads;
+    bool         pipeline;  // true when n_kv_blocks >= FA_MIN_KV_BLOCKS && n_hvx >= 2
+    uint32_t     n_hvx;
 
     // Op parameters
     __fp16       scale;
@@ -171,7 +171,7 @@ struct hmx_fa_context {
     HVX_Vector * vtcm_l_vec;           // Row sum [g_br]
     HVX_Vector * vtcm_s_rowmax;        // Softmax intermediate [g_br]
     HVX_Vector * vtcm_p_rowsum;        // Softmax intermediate [g_br]
-    HVX_Vector * vtcm_row_bufs;        // Per-thread softmax row scratch [n_threads][2][Bc/64]
+    HVX_Vector * vtcm_row_bufs;        // Per-HVX softmax row scratch [n_hvx][2][Bc/64]
     uint8_t *    vtcm_hmx_scales_id;   // HMX output scales (identity)
     uint8_t *    vtcm_hmx_scales_qk;   // HMX output scales (qk_scale)
     __fp16 *     vtcm_mask_buf;        // VTCM mask buffer [Br * m_line], DMA'd per KV block
@@ -260,7 +260,7 @@ static uint32_t fa_scan_active_blocks(
     return n_active;
 }
 
-static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * data) {
+static void flash_attn_ext_f16_thread(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct htp_fa_context * factx = (struct htp_fa_context *) data;
     const struct htp_ops_context * octx = factx->octx;
     const struct htp_tensor * q     = octx->src[0];
@@ -306,15 +306,15 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
     const uint32_t nb3 = dst->nb[3];
 
     // total rows in q
-    const uint32_t dr  = factx->qrows_per_thread;
+    const uint32_t dr  = factx->qrows_per_hvx;
     const uint32_t ir0 = factx->qrow_start + dr * ith;
     const uint32_t ir1 = MIN(ir0 + dr, factx->qrow_start + factx->qrows);
 
     if (ir0 >= ir1) return;
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
 
     const uint32_t DK = nek0;
     const uint32_t DV = nev0;
@@ -774,7 +774,7 @@ typedef struct {
     size_t                  buf_idx;
 } fa_k_int_args_t;
 
-static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data) {
+static void fa_k_interleave_thread(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     fa_k_int_args_t *       args  = (fa_k_int_args_t *) data;
     struct hmx_fa_context * factx = args->factx;
 
@@ -787,7 +787,7 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
         return;
     }
 
-    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
     if (factx->k->type == HTP_TYPE_Q8_0) {
         for (uint32_t r = start; r < end; ++r) {
@@ -801,18 +801,14 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
 }
 
 static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, void * curr_k, uint32_t kv_start, size_t buf_idx) {
-    work_queue_t wp = factx->octx->ctx->work_queue;
+    work_queue_t wp = factx->octx->ctx->hvx_queue;
     uint32_t n = 1;
-    if (factx->n_threads > 1 && kv_rows >= factx->n_threads * 2) {
-        n = factx->n_threads;
+    if (factx->n_hvx > 1 && kv_rows >= factx->n_hvx * 2) {
+        n = factx->n_hvx;
     }
     uint32_t rows_per_t = hex_align_up(hmx_ceil_div(kv_rows, n), 2);
     fa_k_int_args_t args = { factx, kv_rows, src_stride, curr_k, kv_start, rows_per_t, buf_idx };
-    if (n > 1) {
-        work_queue_run(wp, fa_k_interleave_thread, &args, n);
-    } else {
-        fa_k_interleave_thread(1, 0, &args);
-    }
+    work_queue_run(wp, fa_k_interleave_thread, &args, n);
 }
 
 typedef struct {
@@ -826,7 +822,7 @@ typedef struct {
     uint32_t                rows_per_t;
 } fa_v_int_args_t;
 
-static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data) {
+static void fa_v_interleave_thread(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     fa_v_int_args_t *       args  = (fa_v_int_args_t *) data;
     struct hmx_fa_context * factx = args->factx;
 
@@ -841,7 +837,7 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data)
 
     __fp16 * v_tiles_dst = (__fp16 *) args->v_tiles_dst;
 
-    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
     if (factx->v->type == HTP_TYPE_Q8_0) {
         for (uint32_t r = start; r < end; ++r) {
@@ -861,18 +857,14 @@ static void fa_phase_v_interleave(struct hmx_fa_context * factx,
                                   void *                  v_tiles_dst,
                                   size_t                  n_col_tiles,
                                   uint32_t                kv_start) {
-    work_queue_t wp = factx->octx->ctx->work_queue;
+    work_queue_t wp = factx->octx->ctx->hvx_queue;
     uint32_t n = 1;
-    if (factx->n_threads > 1 && kv_rows >= factx->n_threads * 2) {
-        n = factx->n_threads;
+    if (factx->n_hvx > 1 && kv_rows >= factx->n_hvx * 2) {
+        n = factx->n_hvx;
     }
     uint32_t rows_per_t = hex_align_up(hmx_ceil_div(kv_rows, n), 2);
     fa_v_int_args_t args = { factx, kv_rows, src_stride, v_src, v_tiles_dst, n_col_tiles, kv_start, rows_per_t };
-    if (n > 1) {
-        work_queue_run(wp, fa_v_interleave_thread, &args, n);
-    } else {
-        fa_v_interleave_thread(1, 0, &args);
-    }
+    work_queue_run(wp, fa_v_interleave_thread, &args, n);
 }
 
 typedef struct {
@@ -888,7 +880,7 @@ typedef struct {
     atomic_uint               barrier;
 } fa_q_load_args_t;
 
-static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
+static void fa_q_load_thread(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     fa_q_load_args_t *      args  = (fa_q_load_args_t *) data;
     struct hmx_fa_context * factx = args->factx;
 
@@ -902,7 +894,7 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
     const size_t start      = (size_t) i * rows_per_t;
     const size_t end        = hex_smin(start + rows_per_t, factx->g_br);
 
-    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_Q_PREP, (uint16_t) (args->q_start * G + start));
 
     // Parallel initialization of per-block state
@@ -1010,10 +1002,10 @@ static void fa_phase_q_load(struct hmx_fa_context *   factx,
                             uint32_t                  kv_head,
                             uint32_t                  ib3,
                             size_t                    n_rows_g) {
-    work_queue_t wp = factx->octx->ctx->work_queue;
+    work_queue_t wp = factx->octx->ctx->hvx_queue;
     uint32_t n = 1;
-    if (factx->n_threads > 1 && n_rows_g >= (size_t) (factx->n_threads * 2)) {
-        n = factx->n_threads;
+    if (factx->n_hvx > 1 && n_rows_g >= (size_t) (factx->n_hvx * 2)) {
+        n = factx->n_hvx;
     }
     size_t rows_per_t = hex_align_up(hmx_ceil_div(factx->g_br, n), 2);
     const uint32_t n_rows_q = hex_smin(factx->Br, factx->neq1 - q_start);
@@ -1028,11 +1020,7 @@ static void fa_phase_q_load(struct hmx_fa_context *   factx,
     args.n_rows_q = n_rows_q;
     args.q_transposed = q->nb[1] < q->nb[2];
     atomic_init(&args.barrier, n);
-    if (n > 1) {
-        work_queue_run(wp, fa_q_load_thread, &args, n);
-    } else {
-        fa_q_load_thread(1, 0, &args);
-    }
+    work_queue_run(wp, fa_q_load_thread, &args, n);
 }
 
 typedef struct {
@@ -1046,7 +1034,7 @@ typedef struct {
     size_t                    rows_per_t;
 } fa_o_store_args_t;
 
-static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
+static void fa_o_store_thread_f32(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     fa_o_store_args_t *     args  = (fa_o_store_args_t *) data;
     struct hmx_fa_context * factx = args->factx;
 
@@ -1062,7 +1050,7 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
         return;
     }
 
-    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
 
     const struct htp_tensor * dst        = args->dst;
@@ -1113,7 +1101,7 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
 }
 
-static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
+static void fa_o_store_thread_f16(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     fa_o_store_args_t *     args  = (fa_o_store_args_t *) data;
     struct hmx_fa_context * factx = args->factx;
 
@@ -1128,7 +1116,7 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
         return;
     }
 
-    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
 
     const struct htp_tensor * dst        = args->dst;
@@ -1190,19 +1178,15 @@ static void fa_phase_o_store(struct hmx_fa_context *   factx,
                              uint32_t                  kv_head,
                              uint32_t                  ib3,
                              size_t                    n_rows_g) {
-    work_queue_t wp = factx->octx->ctx->work_queue;
+    work_queue_t wp = factx->octx->ctx->hvx_queue;
     uint32_t n = 1;
-    if (factx->n_threads > 1 && n_rows_g >= (size_t) (factx->n_threads * 2)) {
-        n = factx->n_threads;
+    if (factx->n_hvx > 1 && n_rows_g >= (size_t) (factx->n_hvx * 2)) {
+        n = factx->n_hvx;
     }
     size_t rows_per_t = hmx_ceil_div(n_rows_g, n);
     fa_o_store_args_t args = { factx, dst, o_tile_src, q_start, kv_head, ib3, n_rows_g, rows_per_t };
-    worker_callback_t store_fn = factx->is_dst_fp32 ? fa_o_store_thread_f32 : fa_o_store_thread_f16;
-    if (n > 1) {
-        work_queue_run(wp, store_fn, &args, n);
-    } else {
-        store_fn(1, 0, &args);
-    }
+    work_queue_func_t store_fn = factx->is_dst_fp32 ? fa_o_store_thread_f32 : fa_o_store_thread_f16;
+    work_queue_run(wp, store_fn, &args, n);
 }
 
 typedef struct {
@@ -1229,7 +1213,7 @@ typedef struct {
 } fa_softmax_args_t;
 
 static inline void fa_softmax_impl(
-    unsigned int n, unsigned int i, void * data,
+    const htp_thread_t * th, unsigned int n, unsigned int i, void * data,
     const bool has_mask,
     const bool mask_broadcast,
     const bool is_g1,
@@ -1259,7 +1243,7 @@ static inline void fa_softmax_impl(
         return;
     }
 
-    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_SFM, (uint16_t) (args->q_start * G + vec_start * 64));
 
     // Per-thread row scratch: thread i uses bufs at offset i * 2 * stride
@@ -1619,8 +1603,8 @@ static inline void fa_softmax_impl(
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_SFM, (uint16_t) (args->q_start * G + vec_start * 64));
 }
 
-static void fa_softmax_thread_nomask(unsigned int n, unsigned int i, void * data) {
-    fa_softmax_impl(n, i, data,
+static void fa_softmax_thread_nomask(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
+    fa_softmax_impl(th, n, i, data,
                     /*has_mask=*/false,
                     /*mask_broadcast=*/false,
                     /*is_g1=*/false,
@@ -1628,8 +1612,8 @@ static void fa_softmax_thread_nomask(unsigned int n, unsigned int i, void * data
                     /*has_softcap=*/false);
 }
 
-static void fa_softmax_thread_mask_broadcast_g1(unsigned int n, unsigned int i, void * data) {
-    fa_softmax_impl(n, i, data,
+static void fa_softmax_thread_mask_broadcast_g1(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
+    fa_softmax_impl(th, n, i, data,
                     /*has_mask=*/true,
                     /*mask_broadcast=*/true,
                     /*is_g1=*/true,
@@ -1637,8 +1621,8 @@ static void fa_softmax_thread_mask_broadcast_g1(unsigned int n, unsigned int i,
                     /*has_softcap=*/false);
 }
 
-static void fa_softmax_thread_mask_broadcast_gn(unsigned int n, unsigned int i, void * data) {
-    fa_softmax_impl(n, i, data,
+static void fa_softmax_thread_mask_broadcast_gn(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
+    fa_softmax_impl(th, n, i, data,
                     /*has_mask=*/true,
                     /*mask_broadcast=*/true,
                     /*is_g1=*/false,
@@ -1646,7 +1630,7 @@ static void fa_softmax_thread_mask_broadcast_gn(unsigned int n, unsigned int i,
                     /*has_softcap=*/false);
 }
 
-static void fa_softmax_thread(unsigned int n, unsigned int i, void * data) {
+static void fa_softmax_thread(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     fa_softmax_args_t *     args  = (fa_softmax_args_t *) data;
     struct hmx_fa_context * factx = args->factx;
 
@@ -1656,7 +1640,7 @@ static void fa_softmax_thread(unsigned int n, unsigned int i, void * data) {
     const bool has_alibi      = args->has_alibi;
     const bool has_softcap    = factx->has_softcap;
 
-    fa_softmax_impl(n, i, data, has_mask, mask_broadcast, is_g1, has_alibi, has_softcap);
+    fa_softmax_impl(th, n, i, data, has_mask, mask_broadcast, is_g1, has_alibi, has_softcap);
 }
 
 static inline HVX_Vector hvx_fa_load_sinks(const float * sinks_data, size_t idx, size_t count) {
@@ -1742,10 +1726,10 @@ static void fa_phase_softmax_and_build_d(struct hmx_fa_context * factx,
                                          fa_softmax_args_t *     sargs,
                                          size_t                  n_row_tiles,
                                          size_t                  n_row_tiles_g_br) {
-    work_queue_t wp = factx->octx->ctx->work_queue;
+    work_queue_t wp = factx->octx->ctx->hvx_queue;
     const size_t n_row_vec_cnt = hmx_ceil_div(sargs->n_rows_g, 64);
 
-    worker_callback_t softmax_fn = fa_softmax_thread;
+    work_queue_func_t softmax_fn = fa_softmax_thread;
     if (sargs->mask == NULL && !factx->has_softcap && !sargs->has_alibi) {
         softmax_fn = fa_softmax_thread_nomask;
     } else if (sargs->mask != NULL && factx->mask_broadcast && !factx->has_softcap && !sargs->has_alibi) {
@@ -1756,13 +1740,12 @@ static void fa_phase_softmax_and_build_d(struct hmx_fa_context * factx,
         }
     }
 
-    if (factx->n_threads > 1 && n_row_vec_cnt >= 2) {
-        uint32_t n_use = (uint32_t) hex_smin((size_t) factx->n_threads, n_row_vec_cnt);
+    uint32_t n_use = 1;
+    if (factx->n_hvx > 1 && n_row_vec_cnt >= 2) {
+        n_use = (uint32_t) hex_smin((size_t) factx->n_hvx, n_row_vec_cnt);
         sargs->thread_div = init_fastdiv_values(n_use);
-        work_queue_run(wp, softmax_fn, sargs, n_use);
-    } else {
-        softmax_fn(1, 0, sargs);
     }
+    work_queue_run(wp, softmax_fn, sargs, n_use);
 }
 
 // ============================================================================
@@ -1780,7 +1763,7 @@ typedef struct {
     uint8_t *      hmx_scales;
 } hmx_fa_qk_job_t;
 
-static void hmx_fa_qk_dot_worker(void * data) {
+static void hmx_fa_qk_dot_worker(const htp_thread_t * th, void * data) {
     hmx_fa_qk_job_t * job            = (hmx_fa_qk_job_t *) data;
     const size_t      n_row_tiles    = job->n_row_tiles;
     const size_t      n_col_tiles    = job->n_col_tiles;
@@ -1822,7 +1805,7 @@ typedef struct {
     size_t         DV;
 } hmx_fa_o_update_job_t;
 
-static void hmx_fa_o_update_worker(void * data) {
+static void hmx_fa_o_update_worker(const htp_thread_t * th, void * data) {
     hmx_fa_o_update_job_t * job              = (hmx_fa_o_update_job_t *) data;
     const size_t            n_row_tiles      = job->n_row_tiles;
     const size_t            n_col_tiles      = job->n_col_tiles;
@@ -1867,7 +1850,7 @@ typedef struct {
     size_t         DV;
 } hmx_fa_o_norm_job_t;
 
-static void hmx_fa_o_norm_worker(void * data) {
+static void hmx_fa_o_norm_worker(const htp_thread_t * th, void * data) {
     hmx_fa_o_norm_job_t * job              = (hmx_fa_o_norm_job_t *) data;
     const size_t          n_row_tiles      = job->n_row_tiles;
     const size_t          n_row_tiles_g_br = job->n_row_tiles_g_br;
@@ -1992,7 +1975,6 @@ static inline void fa_prefetch_block(dma_queue * dma_q, const struct htp_tensor
 
 int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     struct htp_thread_trace * tr_hvx = &octx->ctx->trace[0];
-    struct htp_thread_trace * tr_hmx = &octx->ctx->trace[HTP_MAX_NTHREADS];
     const struct htp_tensor * q    = octx->src[0];
     const struct htp_tensor * k    = octx->src[1];
     const struct htp_tensor * v    = octx->src[2];
@@ -2037,7 +2019,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     factx.sinks          = octx->src[4];  // NULL if this op has no attention sinks
     factx.k              = k;
     factx.v              = v;
-    factx.n_threads      = kparams->n_threads;
+    factx.n_hvx          = kparams->n_hvx;
     factx.DK             = DK;
     factx.DV             = DV;
     factx.DK_pad         = DK_pad;
@@ -2077,7 +2059,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     const uint32_t Bc = factx.Bc;
     const uint32_t g_br = factx.g_br;
     const bool pipeline = factx.pipeline;
-    const uint32_t n_threads = factx.n_threads;
+    const uint32_t n_hvx = factx.n_hvx;
     const uint32_t G = factx.G;
 
     // Multi-device: prefer head-parallel partitioning (each core owns a disjoint head
@@ -2147,7 +2129,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     // Build the VTCM layout once (shared with the host estimator) and place every
     // scratch buffer at its computed offset. Padded head dims size the HMX tiles.
     struct hmx_fa_vtcm_layout L;
-    hmx_fa_vtcm_layout_build(&L, G, DK_pad, DV_pad, Br, Bc, n_threads, pipeline, factx.is_q_fp32, factx.sinks != NULL, factx.n_heads);
+    hmx_fa_vtcm_layout_build(&L, G, DK_pad, DV_pad, Br, Bc, n_hvx, pipeline, factx.is_q_fp32, factx.sinks != NULL, factx.n_heads);
 
     if (L.total_bytes > ctx->vtcm_size) {
         return HTP_STATUS_VTCM_TOO_SMALL;
@@ -2681,7 +2663,7 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+    if (!htp_ops_context_set_n_hvx(octx, kparams->n_hvx)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
@@ -2776,16 +2758,16 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_hvx = octx->n_hvx;
 
     factx.qrows            = qrows;
     factx.qrow_start       = qrow_start;
-    factx.qrows_per_thread = fastdiv(qrows + n_threads - 1, &octx->n_threads_div);
+    factx.qrows_per_hvx    = fastdiv(qrows + n_hvx - 1, &octx->n_hvx_div);
 
     const bool has_mask  = (mask != NULL);
     const bool has_sinks = (octx->src[4] != NULL);
     struct hvx_fa_vtcm_layout L;
-    hvx_fa_vtcm_layout_build(&L, k->ne[0], v->ne[0], factx.is_q_fp32, has_mask, has_sinks, n_head, n_threads);
+    hvx_fa_vtcm_layout_build(&L, k->ne[0], v->ne[0], factx.is_q_fp32, has_mask, has_sinks, n_head, n_hvx);
 
     if (L.total_bytes > octx->ctx->vtcm_size) {
         return HTP_STATUS_VTCM_TOO_SMALL;
@@ -2822,7 +2804,7 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
         dma_queue_pop(dma_q);
     }
 
-    work_queue_run(octx->ctx->work_queue, flash_attn_ext_f16_thread, &factx, n_threads);
+    work_queue_run(octx->ctx->hvx_queue, flash_attn_ext_f16_thread, &factx, n_hvx);
 
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/flash-attn-ops.h

@@ -46,7 +46,7 @@ struct htp_fa_kernel_params {
     uint8_t  kernel_type;        // enum htp_fa_kernel_type
     uint8_t  head_split;         // 1 = partition by KV heads in multicore, 0 = token partition
     uint8_t  flags;              // HTP_FA_FLAG_*
-    uint8_t  n_threads;          // Number of threads to run
+    uint8_t  n_hvx;              // Number of threads to run
 
     // Common parameters
     uint16_t Br;
@@ -155,7 +155,7 @@ struct hmx_fa_vtcm_layout {
 
 static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
                                             size_t gqa_factor, size_t DK, size_t DV,
-                                            size_t Br, size_t Bc, size_t n_threads,
+                                            size_t Br, size_t Bc, size_t n_hvx,
                                             bool pipeline, bool is_q_fp32,
                                             bool has_sinks, size_t n_heads) {
     const size_t g_br         = hex_align_up(gqa_factor * Br, HMX_FP16_TILE_N_ROWS);
@@ -207,7 +207,7 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
     VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_p_tiles[1], s_tile_size, pipeline);
     VTCM_LAYOUT_ALLOC(off_group_b, off_s_rowmax,      col_vec_size);
     VTCM_LAYOUT_ALLOC(off_group_b, off_p_rowsum,      col_vec_size);
-    VTCM_LAYOUT_ALLOC(off_group_b, off_row_bufs,      row_vec_size * 2 * n_threads);
+    VTCM_LAYOUT_ALLOC(off_group_b, off_row_bufs,      row_vec_size * 2 * n_hvx);
 
     const size_t group_b_size = off_group_b - off_group_b_c;
 
@@ -251,9 +251,9 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
 }
 
 // Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration.
-static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32, bool has_sinks, size_t n_heads) {
+static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_hvx, bool pipeline, bool is_q_fp32, bool has_sinks, size_t n_heads) {
     struct hmx_fa_vtcm_layout L;
-    hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline, is_q_fp32, has_sinks, n_heads);
+    hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_hvx, pipeline, is_q_fp32, has_sinks, n_heads);
     return L.total_bytes;
 }
 
@@ -282,7 +282,7 @@ static inline void hvx_fa_vtcm_layout_build(struct hvx_fa_vtcm_layout * L,
                                             size_t DK, size_t DV,
                                             bool is_q_fp32, bool has_mask,
                                             bool has_sinks, size_t n_heads,
-                                            size_t n_threads) {
+                                            size_t n_hvx) {
     const size_t size_q_row_padded = hex_round_up(DK * (is_q_fp32 ? 4 : 2), 128);
     const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128);
     const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128);
@@ -296,11 +296,11 @@ static inline void hvx_fa_vtcm_layout_build(struct hvx_fa_vtcm_layout * L,
 
     size_t off = 0;
 
-    VTCM_LAYOUT_ALLOC(off, off_q, size_q_block * n_threads);
-    VTCM_LAYOUT_ALLOC(off, off_k, size_k_block * 2 * n_threads);
-    VTCM_LAYOUT_ALLOC(off, off_v, size_v_block * 2 * n_threads);
-    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_m, size_m_block * HVX_FA_DMA_CACHE_SIZE * n_threads, has_mask);
-    VTCM_LAYOUT_ALLOC(off, off_a, size_vkq_acc * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_q, size_q_block * n_hvx);
+    VTCM_LAYOUT_ALLOC(off, off_k, size_k_block * 2 * n_hvx);
+    VTCM_LAYOUT_ALLOC(off, off_v, size_v_block * 2 * n_hvx);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_m, size_m_block * HVX_FA_DMA_CACHE_SIZE * n_hvx, has_mask);
+    VTCM_LAYOUT_ALLOC(off, off_a, size_vkq_acc * n_hvx);
     VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_sinks, size_sinks, has_sinks);
 
     L->size_q_block = size_q_block;
@@ -312,9 +312,9 @@ static inline void hvx_fa_vtcm_layout_build(struct hvx_fa_vtcm_layout * L,
     L->total_bytes  = off;
 }
 
-static inline size_t hvx_fa_compute_vtcm_usage(size_t DK, size_t DV, bool is_q_fp32, bool has_mask, bool has_sinks, size_t n_heads, size_t n_threads) {
+static inline size_t hvx_fa_compute_vtcm_usage(size_t DK, size_t DV, bool is_q_fp32, bool has_mask, bool has_sinks, size_t n_heads, size_t n_hvx) {
     struct hvx_fa_vtcm_layout L;
-    hvx_fa_vtcm_layout_build(&L, DK, DV, is_q_fp32, has_mask, has_sinks, n_heads, n_threads);
+    hvx_fa_vtcm_layout_build(&L, DK, DV, is_q_fp32, has_mask, has_sinks, n_heads, n_hvx);
     return L.total_bytes;
 }
 
@@ -329,7 +329,7 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
                                          size_t   qo_len,
                                          size_t   kv_len,
                                          size_t   vtcm_budget,
-                                         size_t   n_threads,
+                                         size_t   n_hvx,
                                          bool     is_q_fp32,
                                          bool     has_sinks,
                                          size_t   n_heads,
@@ -337,7 +337,7 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
     const size_t T       = HMX_FP16_TILE_N_ROWS;  // 32
     const size_t br_unit = hmx_ceil_div(T, gqa_factor);
     const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2;  // 64
-    const bool   can_pipeline = (kv_len >= FA_MIN_KV_BLOCKS * bc_unit && n_threads >= 2);
+    const bool   can_pipeline = (kv_len >= FA_MIN_KV_BLOCKS * bc_unit && n_hvx >= 2);
 
     // Br_max: largest Br aligned to br_unit that does not exceed qo_len.
     const size_t Br_max = qo_len >= br_unit ? hex_align_down(qo_len, br_unit) : br_unit;
@@ -361,15 +361,15 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
     for (size_t Br = Br_max; Br >= br_unit; Br -= br_unit) {
         // Try all Bc candidates from Bc_limit down to bc_unit
         for (size_t Bc = Bc_limit; Bc >= bc_unit; Bc -= bc_unit) {
-            size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline, is_q_fp32, has_sinks, n_heads);
+            size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_hvx, can_pipeline, is_q_fp32, has_sinks, n_heads);
             if (vtcm_needed <= vtcm_budget) {
                 // This Bc fits for this Br!
                 const size_t q_blocks       = (qo_len + Br - 1) / Br;
                 const size_t kv_blocks      = (kv_len + Bc - 1) / Bc;
-                const size_t actual_threads = (kv_blocks >= 3 && n_threads >= 2) ? n_threads : 1;
+                const size_t actual_hvx     = (kv_blocks >= 3 && n_hvx >= 2) ? n_hvx : 1;
                 const size_t n_rows_g       = Br * gqa_factor;
                 const size_t n_row_vec_cnt  = (n_rows_g + 63) / 64;
-                const size_t n_use          = n_row_vec_cnt < actual_threads ? n_row_vec_cnt : actual_threads;
+                const size_t n_use          = n_row_vec_cnt < actual_hvx ? n_row_vec_cnt : actual_hvx;
                 const size_t vecs_per_t     = n_use > 0 ? (n_row_vec_cnt + n_use - 1) / n_use : 1;
 
                 const size_t c_iter_actual  = c_iter_base + c_softmax * vecs_per_t;

ggml/src/ggml-hexagon/htp/gated-delta-net-ops.c

@@ -748,7 +748,7 @@ static inline void gdn_step_scalar_f32(
     }
 }
 
-static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, void * data) {
+static void gated_delta_net_f32_pp_thread(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct htp_gdn_context * gctx = (struct htp_gdn_context *) data;
     struct htp_ops_context * octx = gctx->octx;
     const struct htp_gdn_kernel_params * kparams = gctx->kparams;
@@ -777,10 +777,10 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
     float * dst_base       = (float *) (uintptr_t) dst->data;
     float * state_out_base = dst_cache ? (float *) (uintptr_t) dst_cache->data : (dst_base + S_v * H * n_tokens * n_seqs);
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     const struct htp_gdn_vtcm_layout * layout = &gctx->layout;
     float * s_work[2];
-    s_work[0] = (float *) (gctx->vtcm_base + layout->bytes_per_thread * ith);
+    s_work[0] = (float *) (gctx->vtcm_base + layout->bytes_per_hvx * ith);
     s_work[1] = s_work[0] + layout->state_aligned / sizeof(float);
 
     const struct fastdiv_values * fd_H   = &kparams->div_H;
@@ -817,7 +817,7 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
         spad_idx ^= 1;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     int curr_spad_idx = 0;
     for (uint32_t ir = gctx->row_start + ith; ir < row_end; ir += nth) {
@@ -892,7 +892,7 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
     dma_queue_flush(dma_q);
 }
 
-static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, void * data) {
+static void gated_delta_net_f32_tg_thread(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct htp_gdn_context * gctx = (struct htp_gdn_context *) data;
     struct htp_ops_context * octx = gctx->octx;
     const struct htp_gdn_kernel_params * kparams = gctx->kparams;
@@ -918,10 +918,10 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo
     const float scale = kparams->scale;
     float * dst_base  = (float *) (uintptr_t) dst->data;
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
     const struct htp_gdn_vtcm_layout * layout = &gctx->layout;
     float * s_work[2];
-    s_work[0] = (float *) (gctx->vtcm_base + layout->bytes_per_thread * ith);
+    s_work[0] = (float *) (gctx->vtcm_base + layout->bytes_per_hvx * ith);
     s_work[1] = s_work[0] + layout->state_aligned / sizeof(float);
 
     const struct fastdiv_values * fd_H   = &kparams->div_H;
@@ -957,7 +957,7 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo
         spad_idx ^= 1;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     int curr_spad_idx = 0;
     for (uint32_t ir = gctx->row_start + ith; ir < row_end; ir += nth) {
@@ -1031,7 +1031,7 @@ struct htp_gdn_hmx_gemm_task {
     uint8_t *      hmx_scales;
 };
 
-static void htp_gdn_hmx_gemm_worker(void * data) {
+static void htp_gdn_hmx_gemm_worker(const htp_thread_t * th, void * data) {
     struct htp_gdn_hmx_gemm_task * task = (struct htp_gdn_hmx_gemm_task *) data;
     asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)task->hmx_scales));
 
@@ -1623,7 +1623,7 @@ static inline void gdn_init_head_ptrs(
     head->rows_inv = VTCM_LAYOUT_PTR(HVX_Vector, vtcm_base, L->off_rows_inv + h * row_vecs_sz);
     head->rows_a   = VTCM_LAYOUT_PTR(HVX_Vector, vtcm_base, L->off_rows_a + h * row_vecs_sz);
 
-    head->vtcm_m   = VTCM_LAYOUT_PTR(HVX_Vector, vtcm_base, L->off_thread_scratch + h * (64 * 128));
+    head->vtcm_m   = VTCM_LAYOUT_PTR(HVX_Vector, vtcm_base, L->off_hvx_scratch + h * (64 * 128));
     head->vtcm_tmp = head->vtcm_m + 32;
 
     head->iv1 = base_iv1 + h;
@@ -1654,10 +1654,10 @@ struct htp_gdn_batch_context {
     const struct htp_gdn_kernel_params * kparams;
 };
 
-static void gdn_hvx_init_state_worker(unsigned int n, unsigned int i, void * data) {
+static void gdn_hvx_init_state_worker(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     (void) n;
     struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data;
-    struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
 
     struct htp_gdn_head_ptrs * head = &bctx->heads[i];
@@ -1690,10 +1690,10 @@ static inline HVX_Vector hvx_splat_last_f32(HVX_Vector v) {
     return hvx_vec_repl4(Q6_V_vror_VR(v, 124));
 }
 
-static void gdn_hvx_phase1a_worker(unsigned int n, unsigned int i, void * data) {
+static void gdn_hvx_phase1a_worker(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     (void) n;
     struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data;
-    struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     const uint16_t info = (uint16_t) bctx->c;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_PREP, info);
 
@@ -1763,10 +1763,10 @@ static void gdn_hvx_phase1a_worker(unsigned int n, unsigned int i, void * data)
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_PREP, info);
 }
 
-static void gdn_hvx_phase1b_worker(unsigned int n, unsigned int i, void * data) {
+static void gdn_hvx_phase1b_worker(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     (void) n;
     struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data;
-    struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     const uint16_t info = (uint16_t) bctx->c;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_PREP, info);
 
@@ -1860,10 +1860,10 @@ static void gdn_hvx_phase1b_worker(unsigned int n, unsigned int i, void * data)
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_PREP, info);
 }
 
-static void gdn_hvx_phase2_worker(unsigned int n, unsigned int i, void * data) {
+static void gdn_hvx_phase2_worker(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     (void) n;
     struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data;
-    struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     const uint16_t info = (uint16_t) bctx->c;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_SOLVE, info);
 
@@ -1884,10 +1884,10 @@ static void gdn_hvx_phase2_worker(unsigned int n, unsigned int i, void * data) {
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_SOLVE, info);
 }
 
-static void gdn_hvx_phase3_worker(unsigned int n, unsigned int i, void * data) {
+static void gdn_hvx_phase3_worker(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     (void) n;
     struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data;
-    struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     const uint16_t info = (uint16_t) bctx->c;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_V_PREP, info);
 
@@ -1928,10 +1928,10 @@ static void gdn_hvx_phase3_worker(unsigned int n, unsigned int i, void * data) {
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_V_PREP, info);
 }
 
-static void gdn_hvx_phase4_worker(unsigned int n, unsigned int i, void * data) {
+static void gdn_hvx_phase4_worker(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     (void) n;
     struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data;
-    struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     const uint16_t info = (uint16_t) bctx->c;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_D_PREP, info);
 
@@ -1959,10 +1959,10 @@ static void gdn_hvx_phase4_worker(unsigned int n, unsigned int i, void * data) {
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_D_PREP, info);
 }
 
-static void gdn_hvx_phase5_worker(unsigned int n, unsigned int i, void * data) {
+static void gdn_hvx_phase5_worker(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     (void) n;
     struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data;
-    struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     const uint16_t info = (uint16_t) bctx->c;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_OUT, info);
 
@@ -1984,10 +1984,10 @@ static void gdn_hvx_phase5_worker(unsigned int n, unsigned int i, void * data) {
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_OUT, info);
 }
 
-static void gdn_hvx_phase6_worker(unsigned int n, unsigned int i, void * data) {
+static void gdn_hvx_phase6_worker(const htp_thread_t * th, unsigned int n, unsigned int i, void * data) {
     (void) n;
     struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data;
-    struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i];
+    struct htp_thread_trace * tr = th->trace;
     const uint16_t info = (uint16_t) bctx->c;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_STATE, info);
 
@@ -2040,7 +2040,7 @@ static int gated_delta_net_f32_hmx_chunked(
     const uint32_t n_sv_tiles = S_v / 32;
 
     struct htp_gdn_hmx_vtcm_layout L;
-    htp_gdn_hmx_vtcm_layout_build(&L, S_v, chunk_size, kparams->n_heads_batch, kparams->n_threads, kparams->pipeline != 0);
+    htp_gdn_hmx_vtcm_layout_build(&L, S_v, chunk_size, kparams->n_heads_batch, kparams->n_hvx, kparams->pipeline != 0);
 
     if (L.total_bytes > octx->ctx->vtcm_size) {
         return HTP_STATUS_VTCM_TOO_SMALL;
@@ -2062,7 +2062,7 @@ static int gated_delta_net_f32_hmx_chunked(
 
     hmx_queue_t  hmx_q = octx->ctx->hmx_queue;
     dma_queue *  dma_q = octx->ctx->dma[0];
-    work_queue_t wp    = octx->ctx->work_queue;
+    work_queue_t wp    = octx->ctx->hvx_queue;
 
     struct htp_gdn_head_ptrs heads[8];
     struct htp_gdn_hmx_gemm_task gemm_tasks[8][9];
@@ -2286,7 +2286,7 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
     }
 
     const struct htp_gdn_kernel_params * kparams = (const struct htp_gdn_kernel_params *) octx->kernel_params;
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+    if (!htp_ops_context_set_n_hvx(octx, kparams->n_hvx)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
@@ -2321,7 +2321,7 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
         return gated_delta_net_f32_hmx_chunked(octx, kparams, row_start, nrows);
     }
 
-    const uint32_t n_threads = (nrows < octx->n_threads) ? nrows : octx->n_threads;
+    const uint32_t n_hvx = (nrows < octx->n_hvx) ? nrows : octx->n_hvx;
 
     struct htp_gdn_context gctx;
     gctx.octx      = octx;
@@ -2329,7 +2329,7 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
     gctx.row_start = row_start;
     gctx.nrows     = nrows;
 
-    htp_gdn_vtcm_layout_build(&gctx.layout, S_v, n_threads);
+    htp_gdn_vtcm_layout_build(&gctx.layout, S_v, n_hvx);
     gctx.vtcm_base = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, gctx.layout.off_state);
 
     if (gctx.layout.total_bytes > octx->ctx->vtcm_size) {
@@ -2337,18 +2337,18 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
     }
 
     FARF(HIGH, "gated-delta-net-f32: q(%ux%ux%ux%u) k(%ux%ux%ux%u) v(%ux%ux%ux%u) state(%ux%ux%ux%u) -> (%ux%ux%ux%u) : "
-         "vtcm-size %zu n_threads %u\n",
+         "vtcm-size %zu n_hvx %u\n",
          octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
          octx->src[1]->ne[0], octx->src[1]->ne[1], octx->src[1]->ne[2], octx->src[1]->ne[3],
          octx->src[2]->ne[0], octx->src[2]->ne[1], octx->src[2]->ne[2], octx->src[2]->ne[3],
          octx->src[5]->ne[0], octx->src[5]->ne[1], octx->src[5]->ne[2], octx->src[5]->ne[3],
          dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         gctx.layout.total_bytes, n_threads);
+         gctx.layout.total_bytes, n_hvx);
 
     if (n_tokens == 1) {
-        work_queue_run(octx->ctx->work_queue, gated_delta_net_f32_tg_thread, &gctx, n_threads);
+        work_queue_run(octx->ctx->hvx_queue, gated_delta_net_f32_tg_thread, &gctx, n_hvx);
     } else {
-        work_queue_run(octx->ctx->work_queue, gated_delta_net_f32_pp_thread, &gctx, n_threads);
+        work_queue_run(octx->ctx->hvx_queue, gated_delta_net_f32_pp_thread, &gctx, n_hvx);
     }
 
     return HTP_STATUS_OK;

ggml/src/ggml-hexagon/htp/gated-delta-net-ops.h

@@ -29,7 +29,7 @@ struct htp_gdn_kernel_params {
     uint16_t n_chunks;
     uint16_t n_heads_batch;
 
-    uint32_t n_threads;
+    uint32_t n_hvx;
     uint32_t S_v;
     uint32_t H;
     uint32_t n_tokens;
@@ -59,7 +59,7 @@ _Static_assert(sizeof(struct htp_gdn_kernel_params) <= 128, "htp_gdn_kernel_para
 
 struct htp_gdn_vtcm_layout {
     size_t state_aligned;
-    size_t bytes_per_thread;
+    size_t bytes_per_hvx;
     size_t off_state;
     size_t total_bytes;
 };
@@ -67,13 +67,13 @@ struct htp_gdn_vtcm_layout {
 static inline void htp_gdn_vtcm_layout_build(
     struct htp_gdn_vtcm_layout * L,
     uint32_t S_v,
-    uint32_t n_threads
+    uint32_t n_hvx
 ) {
     size_t state_bytes = (size_t) S_v * S_v * sizeof(float);
     L->state_aligned = hex_round_up(state_bytes, 128);
-    L->bytes_per_thread = 2 * L->state_aligned;
+    L->bytes_per_hvx = 2 * L->state_aligned;
     size_t off = 0;
-    VTCM_LAYOUT_ALLOC(off, off_state, L->bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_state, L->bytes_per_hvx * n_hvx);
     L->total_bytes = off;
 }
 
@@ -128,7 +128,7 @@ struct htp_gdn_hmx_vtcm_layout {
     size_t off_rows_inv;
     size_t off_rows_a;
 
-    size_t off_thread_scratch;
+    size_t off_hvx_scratch;
     size_t off_scales_1;
 
     size_t state_f32_bytes;
@@ -140,7 +140,7 @@ struct htp_gdn_hmx_vtcm_layout {
     size_t tile_64x64_bytes;
 
     uint32_t n_heads_batch;
-    uint32_t n_threads;
+    uint32_t n_hvx;
     bool     pipeline;
     size_t   total_bytes;
 };
@@ -150,16 +150,16 @@ static inline void htp_gdn_hmx_vtcm_layout_build(
     uint32_t S_v,
     uint32_t chunk_size,
     uint32_t n_heads_batch,
-    uint32_t n_threads,
+    uint32_t n_hvx,
     bool     pipeline
 ) {
     memset(L, 0, sizeof(*L));
     L->n_heads_batch = n_heads_batch;
-    L->n_threads     = n_threads;
+    L->n_hvx         = n_hvx;
     L->pipeline      = pipeline;
 
     const size_t bh = (size_t) n_heads_batch;
-    const size_t nth = (size_t) (n_threads > 0 ? n_threads : 1);
+    const size_t nth = (size_t) (n_hvx > 0 ? n_hvx : 1);
 
     const size_t state_f32_sz   = hex_round_up(S_v * S_v * sizeof(float), 2048);
     const size_t state_f16_sz   = hex_round_up(S_v * S_v * sizeof(__fp16), 2048);
@@ -248,9 +248,9 @@ static inline void htp_gdn_hmx_vtcm_layout_build(
     VTCM_LAYOUT_ALLOC(off, off_rows_inv,    bh * row_vecs_sz);
     VTCM_LAYOUT_ALLOC(off, off_rows_a,      bh * row_vecs_sz);
 
-    const size_t thread_scratch_sz = 64 * 128;
+    const size_t hvx_scratch_sz = 64 * 128;
     off = hex_align_up(off, HMX_FP16_TILE_SIZE);
-    VTCM_LAYOUT_ALLOC(off, off_thread_scratch, nth * thread_scratch_sz);
+    VTCM_LAYOUT_ALLOC(off, off_hvx_scratch, nth * hvx_scratch_sz);
     off = hex_align_up(off, HMX_FP16_TILE_SIZE);
     VTCM_LAYOUT_ALLOC(off, off_scales_1,       HMX_FP16_TILE_SIZE);
 
@@ -263,7 +263,7 @@ static inline bool htp_gdn_hmx_solve_layout(
     uint32_t chunk_size,
     uint32_t total_rows,
     size_t   vtcm_budget,
-    uint32_t n_threads,
+    uint32_t n_hvx,
     bool     pipeline,
     uint32_t * n_heads_batch_out
 ) {
@@ -271,8 +271,8 @@ static inline bool htp_gdn_hmx_solve_layout(
     if (max_batch > total_rows) {
         max_batch = total_rows;
     }
-    if (max_batch > n_threads) {
-        max_batch = n_threads;
+    if (max_batch > n_hvx) {
+        max_batch = n_hvx;
     }
     static const uint32_t candidates[] = { 8, 6, 4, 2, 1 };
     for (size_t i = 0; i < sizeof(candidates) / sizeof(candidates[0]); ++i) {
@@ -281,7 +281,7 @@ static inline bool htp_gdn_hmx_solve_layout(
             continue;
         }
         struct htp_gdn_hmx_vtcm_layout L;
-        htp_gdn_hmx_vtcm_layout_build(&L, S_v, chunk_size, bh, n_threads, pipeline);
+        htp_gdn_hmx_vtcm_layout_build(&L, S_v, chunk_size, bh, n_hvx, pipeline);
         if (L.total_bytes <= vtcm_budget) {
             *layout_out = L;
             *n_heads_batch_out = bh;
@@ -289,7 +289,7 @@ static inline bool htp_gdn_hmx_solve_layout(
         }
     }
     if (pipeline) {
-        return htp_gdn_hmx_solve_layout(layout_out, S_v, chunk_size, total_rows, vtcm_budget, n_threads, false, n_heads_batch_out);
+        return htp_gdn_hmx_solve_layout(layout_out, S_v, chunk_size, total_rows, vtcm_budget, n_hvx, false, n_heads_batch_out);
     }
     return false;
 }

ggml/src/ggml-hexagon/htp/gather-rows.c

@@ -3,7 +3,7 @@
 #include <assert.h>
 
 #include "hex-common.h"
-#include "hex-profile.h"
+#include "htp-trace.h"
 #include "hvx-utils.h"
 #include "dma-queue.h"
 #include "htp-ctx.h"
@@ -23,7 +23,7 @@ struct gather_rows_task {
     dma_addr_t                            dst;
     uint32_t                              r_start;
     uint32_t                              rows;
-    uint32_t                              rows_per_thread;
+    uint32_t                              rows_per_hvx;
     const struct htp_gather_rows_params * kparams;
     gather_rows_compute_fn_t              compute;
 };
@@ -190,21 +190,21 @@ static void gather_rows_compute_f16_vgn(
     }
 }
 
-static void gather_rows_thread_f32(unsigned int nth, unsigned int ith, void * data) {
+static void gather_rows_thread_f32(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     (void) nth;
     const struct gather_rows_task * task = (const struct gather_rows_task *) data;
     const struct htp_gather_rows_params * kparams = task->kparams;
     const struct htp_tensor * a = task->a;
     const struct htp_tensor * b = task->b;
 
-    const uint32_t r_offset = MIN(ith * task->rows_per_thread, task->rows);
+    const uint32_t r_offset = MIN(ith * task->rows_per_hvx, task->rows);
     const uint32_t r_beg    = task->r_start + r_offset;
-    const uint32_t r_end    = task->r_start + MIN(r_offset + task->rows_per_thread, task->rows);
+    const uint32_t r_end    = task->r_start + MIN(r_offset + task->rows_per_hvx, task->rows);
     if (r_beg >= r_end) {
         return;
     }
 
-    struct htp_thread_trace * tr = &task->octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     struct htp_gather_rows_vtcm_layout layout;
     htp_gather_rows_vtcm_layout_build(&layout, kparams->vg, kparams->region_size, kparams->nr_max, kparams->ne, 4);
@@ -227,7 +227,7 @@ static void gather_rows_thread_f32(unsigned int nth, unsigned int ith, void * da
     gather_rows_init_tables_f32(tab, inc, kparams, a, b);
     htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, (uint16_t) r_beg);
 
-    dma_queue * q = task->octx->ctx->dma[ith];
+    dma_queue * q = th->dma_q;
     const uint32_t n_in   = 1 + (task->b ? (kparams->b_dense ? 1 : kparams->nb) : 0);
     const uint32_t nr_max = kparams->nr_max;
     const uint32_t nr0    = MIN(nr_max, r_end - r_beg);
@@ -271,21 +271,21 @@ static void gather_rows_thread_f32(unsigned int nth, unsigned int ith, void * da
     dma_queue_flush(q);
 }
 
-static void gather_rows_thread_f16(unsigned int nth, unsigned int ith, void * data) {
+static void gather_rows_thread_f16(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     (void) nth;
     const struct gather_rows_task * task = (const struct gather_rows_task *) data;
     const struct htp_gather_rows_params * kparams = task->kparams;
     const struct htp_tensor * a = task->a;
     const struct htp_tensor * b = task->b;
 
-    const uint32_t r_offset = MIN(ith * task->rows_per_thread, task->rows);
+    const uint32_t r_offset = MIN(ith * task->rows_per_hvx, task->rows);
     const uint32_t r_beg    = task->r_start + r_offset;
-    const uint32_t r_end    = task->r_start + MIN(r_offset + task->rows_per_thread, task->rows);
+    const uint32_t r_end    = task->r_start + MIN(r_offset + task->rows_per_hvx, task->rows);
     if (r_beg >= r_end) {
         return;
     }
 
-    struct htp_thread_trace * tr = &task->octx->ctx->trace[ith];
+    struct htp_thread_trace * tr = th->trace;
 
     struct htp_gather_rows_vtcm_layout layout;
     htp_gather_rows_vtcm_layout_build(&layout, kparams->vg, kparams->region_size, kparams->nr_max, kparams->ne, 2);
@@ -308,7 +308,7 @@ static void gather_rows_thread_f16(unsigned int nth, unsigned int ith, void * da
     gather_rows_init_tables_f16(tab, inc, kparams, a, b);
     htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, (uint16_t) r_beg);
 
-    dma_queue * q = task->octx->ctx->dma[ith];
+    dma_queue * q = th->dma_q;
     const uint32_t n_in = 1 + (task->b ? (kparams->b_dense ? 1 : kparams->nb) : 0);
 
     const uint32_t nr_max = kparams->nr_max;
@@ -357,13 +357,13 @@ int htp_gather_rows(
     const struct htp_tensor * a,
     const struct htp_tensor * b,
     const struct htp_gather_rows_params * kparams,
-    uint32_t n_threads) {
+    uint32_t n_hvx) {
 
-    if (!htp_ops_context_set_n_threads(octx, n_threads)) {
+    if (!htp_ops_context_set_n_hvx(octx, n_hvx)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
-    assert(kparams->slice * octx->n_threads <= octx->ctx->vtcm_size);
+    assert(kparams->slice * octx->n_hvx <= octx->ctx->vtcm_size);
 
     const uint32_t total_rows = kparams->total_rows;
     uint32_t r_start = 0;
@@ -389,7 +389,7 @@ int htp_gather_rows(
     task.r_start         = r_start;
     task.rows            = r_count;
     task.kparams         = kparams;
-    task.rows_per_thread = kparams->rows_per_thread;
+    task.rows_per_hvx    = kparams->rows_per_hvx;
 
     if (elem_size == 2) {
         task.compute = (kparams->vg == 1) ? gather_rows_compute_f16_vg1 : gather_rows_compute_f16_vgn;
@@ -398,6 +398,6 @@ int htp_gather_rows(
     }
 
     work_queue_func_t worker_func = (elem_size == 2) ? gather_rows_thread_f16 : gather_rows_thread_f32;
-    work_queue_run(octx->ctx->work_queue, worker_func, &task, octx->n_threads);
+    work_queue_run(octx->ctx->hvx_queue, worker_func, &task, octx->n_hvx);
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/gather-rows.h

@@ -22,7 +22,7 @@ struct htp_gather_rows_params {
     uint32_t region_size;
     uint32_t slice;
     uint32_t total_rows;
-    uint32_t rows_per_thread;
+    uint32_t rows_per_hvx;
     uint8_t  b_dense;
     uint8_t  pad[3];
 };
@@ -35,7 +35,7 @@ struct htp_gather_rows_vtcm_layout {
     size_t off_out0;
     size_t off_region1;
     size_t off_out1;
-    size_t total_per_thread;
+    size_t total_per_hvx;
 };
 
 static inline void htp_gather_rows_vtcm_layout_build(
@@ -58,7 +58,7 @@ static inline void htp_gather_rows_vtcm_layout_build(
     VTCM_LAYOUT_ALLOC(off, off_out0,    out_aligned);
     VTCM_LAYOUT_ALLOC(off, off_region1, region_aligned);
     VTCM_LAYOUT_ALLOC(off, off_out1,    out_aligned);
-    L->total_per_thread = off;
+    L->total_per_hvx = off;
 }
 
 static inline bool htp_gather_rows_solve_layout(
@@ -67,7 +67,7 @@ static inline bool htp_gather_rows_solve_layout(
     uint32_t nb, uint32_t b_nb1, bool b_dense,
     uint32_t elem_size,
     uint32_t total_rows,
-    uint32_t n_threads,
+    uint32_t n_hvx,
     uint32_t mdev_count,
     uint32_t vtcm_size) {
 
@@ -86,13 +86,13 @@ static inline bool htp_gather_rows_solve_layout(
     const uint32_t rg      = v_lanes / g;
     const uint32_t vg      = ne / g;
 
-    const uint32_t slice = (vtcm_size / n_threads) & ~127u;
+    const uint32_t slice = (vtcm_size / n_hvx) & ~127u;
     const uint32_t total_chunks = total_rows / rg;
     const uint32_t chunks_per_dev = (total_chunks + mdev_count - 1) / mdev_count;
     const uint32_t max_rows_per_dev = (mdev_count > 1 && total_chunks >= mdev_count)
         ? (chunks_per_dev * rg + (total_rows % rg))
         : total_rows;
-    const uint32_t rows_per_thread = hex_round_up((max_rows_per_dev + n_threads - 1) / n_threads, rg);
+    const uint32_t rows_per_hvx = hex_round_up((max_rows_per_dev + n_hvx - 1) / n_hvx, rg);
 
     const uint32_t fixed = 3 * vg * 128;
     const uint32_t per_row = a_nb1 + (nb > 0 ? (b_dense ? b_nb1 : nb * b_nb1) : 0) + ne * elem_size;
@@ -101,7 +101,7 @@ static inline bool htp_gather_rows_solve_layout(
     }
 
     const uint32_t avail_per_buf = (slice - fixed) / 2;
-    uint32_t nr_max = MIN(avail_per_buf / per_row, rows_per_thread);
+    uint32_t nr_max = MIN(avail_per_buf / per_row, rows_per_hvx);
     if (elem_size == 2) {
         const uint32_t per_row_region = a_nb1 + (nb > 0 ? (b_dense ? b_nb1 : nb * b_nb1) : 0);
         if (per_row_region > 0) {
@@ -132,7 +132,7 @@ static inline bool htp_gather_rows_solve_layout(
         struct htp_gather_rows_vtcm_layout layout;
         htp_gather_rows_vtcm_layout_build(&layout, vg, region_size, nr_max, ne, elem_size);
 
-        if ((elem_size != 2 || region_size <= 32768) && layout.total_per_thread <= slice) {
+        if ((elem_size != 2 || region_size <= 32768) && layout.total_per_hvx <= slice) {
             break;
         }
 
@@ -154,7 +154,7 @@ static inline bool htp_gather_rows_solve_layout(
     kparams->region_size     = region_size;
     kparams->slice           = slice;
     kparams->total_rows      = total_rows;
-    kparams->rows_per_thread = rows_per_thread;
+    kparams->rows_per_hvx    = rows_per_hvx;
     kparams->b_dense         = b_dense ? 1 : 0;
     kparams->pad[0]          = 0;
     kparams->pad[1]          = 0;
@@ -183,7 +183,7 @@ int htp_gather_rows(
     const struct htp_tensor * a,
     const struct htp_tensor * b,
     const struct htp_gather_rows_params * kparams,
-    uint32_t n_threads);
+    uint32_t n_hvx);
 #endif
 
 #endif // HTP_GATHER_ROWS_H

ggml/src/ggml-hexagon/htp/get-rows-ops.c

@@ -37,7 +37,7 @@ struct get_rows_context {
     get_rows_tiled_compute_fn_t compute_tiled;
     uint32_t task_start;
     uint32_t tasks;
-    uint32_t tasks_per_thread;
+    uint32_t tasks_per_hvx;
     uint32_t tile_size;
     uint32_t tile_stride;
     uint32_t src0_row_size;
@@ -95,20 +95,20 @@ static __attribute__((noinline)) uint32_t get_rows_read_index_vtcm(
     return (uint32_t) Q6_R_vextract_VR(v, off);
 }
 
-static void get_rows_thread_sametype(unsigned int nth, unsigned int ith, void * data) {
+static void get_rows_thread_sametype(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct get_rows_context * grctx = (struct get_rows_context *) data;
     struct htp_ops_context * octx = grctx->octx;
     const struct htp_get_rows_kernel_params * kparams = grctx->kparams;
     get_rows_preamble;
 
-    const uint32_t dr  = grctx->tasks_per_thread;
+    const uint32_t dr  = grctx->tasks_per_hvx;
     const uint32_t ir0 = grctx->task_start + dr * ith;
     if (ir0 >= grctx->task_start + grctx->tasks) {
         return;
     }
     const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);
     const uint32_t row_size_bytes = grctx->src0_row_size;
-    dma_queue * dma_q = octx->ctx->dma[ith];
+    dma_queue * dma_q = th->dma_q;
 
     for (uint32_t i = ir0; i < ir1; ++i) {
         const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11);
@@ -174,23 +174,23 @@ static inline struct get_rows_flat_task get_rows_flat_calc_task(
     return task;
 }
 
-static void get_rows_thread_flat(unsigned int nth, unsigned int ith, void * data) {
+static void get_rows_thread_flat(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct get_rows_context * grctx = (struct get_rows_context *) data;
     struct htp_ops_context * octx = grctx->octx;
     get_rows_preamble;
 
-    const uint32_t dr  = grctx->tasks_per_thread;
+    const uint32_t dr  = grctx->tasks_per_hvx;
     const uint32_t ir0 = grctx->task_start + dr * ith;
     if (ir0 >= grctx->task_start + grctx->tasks) {
         return;
     }
     const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);
 
-    dma_queue * dma_q = octx->ctx->dma[ith];
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    dma_queue * dma_q = th->dma_q;
+    struct htp_thread_trace * tr = th->trace;
     const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout;
-    uint8_t * vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_src0) + ith * vtcm_layout->src0_bytes_per_thread;
-    uint8_t * vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_dst)  + ith * vtcm_layout->dst_bytes_per_thread;
+    uint8_t * vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_src0) + ith * vtcm_layout->src0_bytes_per_hvx;
+    uint8_t * vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_dst)  + ith * vtcm_layout->dst_bytes_per_hvx;
     const uint32_t src0_bytes = grctx->src0_row_size;
     const uint32_t dst_bytes  = grctx->dst_row_size;
     get_rows_flat_compute_fn_t compute = grctx->compute_flat;
@@ -276,12 +276,12 @@ static inline struct get_rows_tiled_task get_rows_tiled_calc_task(
     return task;
 }
 
-static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * data) {
+static void get_rows_thread_tiled(const htp_thread_t * th, unsigned int nth, unsigned int ith, void * data) {
     struct get_rows_context * grctx = (struct get_rows_context *) data;
     struct htp_ops_context * octx = grctx->octx;
     get_rows_preamble;
 
-    const uint32_t dr  = grctx->tasks_per_thread;
+    const uint32_t dr  = grctx->tasks_per_hvx;
     const uint32_t ir0 = grctx->task_start + dr * ith;
     if (ir0 >= grctx->task_start + grctx->tasks) {
         return;
@@ -290,10 +290,10 @@ static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * dat
     const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);
     const uint32_t n_k_tiles = ne00 / htp_tensor_get_tile_k(octx->src[0]->type);
     const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout;
-    uint8_t * src_spad_base = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_src0) + ith * vtcm_layout->src0_bytes_per_thread;
-    uint8_t * dst_spad_base = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_dst)  + ith * vtcm_layout->dst_bytes_per_thread;
-    dma_queue * dma_q = octx->ctx->dma[ith];
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    uint8_t * src_spad_base = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_src0) + ith * vtcm_layout->src0_bytes_per_hvx;
+    uint8_t * dst_spad_base = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_dst)  + ith * vtcm_layout->dst_bytes_per_hvx;
+    dma_queue * dma_q = th->dma_q;
+    struct htp_thread_trace * tr = th->trace;
 
     const uint32_t tile_size   = grctx->tile_size;
     const uint32_t tile_stride = grctx->tile_stride;
@@ -414,11 +414,11 @@ int op_get_rows(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+    if (!htp_ops_context_set_n_hvx(octx, kparams->n_hvx)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_hvx = octx->n_hvx;
 
     struct get_rows_context grctx;
     grctx.octx = octx;
@@ -428,14 +428,14 @@ int op_get_rows(struct htp_ops_context * octx) {
     grctx.compute_tiled = NULL;
     grctx.task_start = task_start;
     grctx.tasks = tasks;
-    grctx.tasks_per_thread = fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
+    grctx.tasks_per_hvx = fastdiv(tasks + n_hvx - 1, &octx->n_hvx_div);
     grctx.tile_size   = 0;
     grctx.tile_stride = 0;
     grctx.src0_row_size = src0_row_size;
     grctx.dst_row_size  = dst_row_size;
     grctx.idx_type      = octx->src[1]->type;
 
-    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, kparams->kernel_type, octx->src[0]->type, ne00, total_tasks, grctx.idx_type, n_threads);
+    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, kparams->kernel_type, octx->src[0]->type, ne00, total_tasks, grctx.idx_type, n_hvx);
 
     if (grctx.vtcm_layout.total_bytes > octx->ctx->vtcm_size) {
         FARF(ERROR, "get-rows: VTCM reservation %zu is too small, needed %zu\n",
@@ -497,14 +497,14 @@ int op_get_rows(struct htp_ops_context * octx) {
         grctx.read_index = get_rows_read_index_vtcm;
     }
 
-    FARF(HIGH, "get-rows: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu dst-vtcm-size %zu kernel-type %d n-threads %d\n",
+    FARF(HIGH, "get-rows: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu dst-vtcm-size %zu kernel-type %d n-hvx %d\n",
          octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
          octx->src[1]->ne[0], octx->src[1]->ne[1], octx->src[1]->ne[2], octx->src[1]->ne[3],
          octx->dst->ne[0], octx->dst->ne[1], octx->dst->ne[2], octx->dst->ne[3],
-         grctx.vtcm_layout.src0_bytes_per_thread * n_threads,
-         grctx.vtcm_layout.dst_bytes_per_thread  * n_threads,
-         kparams->kernel_type, n_threads);
+         grctx.vtcm_layout.src0_bytes_per_hvx * n_hvx,
+         grctx.vtcm_layout.dst_bytes_per_hvx  * n_hvx,
+         kparams->kernel_type, n_hvx);
 
-    work_queue_run(octx->ctx->work_queue, q_func, &grctx, n_threads);
+    work_queue_run(octx->ctx->hvx_queue, q_func, &grctx, n_hvx);
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/get-rows-ops.h

@@ -16,7 +16,7 @@ enum htp_get_rows_kernel_type {
 
 struct htp_get_rows_kernel_params {
     uint8_t  kernel_type;
-    uint32_t n_threads;
+    uint32_t n_hvx;
     uint32_t total_tasks;
     uint32_t vtcm_size;
 
@@ -34,8 +34,8 @@ struct htp_get_rows_vtcm_layout {
     size_t off_src0;
     size_t off_dst;
 
-    size_t src0_bytes_per_thread;
-    size_t dst_bytes_per_thread;
+    size_t src0_bytes_per_hvx;
+    size_t dst_bytes_per_hvx;
 
     size_t src0_spad_half_size;
     size_t dst_spad_half_size;
@@ -48,7 +48,7 @@ static inline void htp_get_rows_vtcm_layout_build(
     uint32_t ne00,
     uint32_t total_tasks,
     int idx_type,
-    uint32_t n_threads) {
+    uint32_t n_hvx) {
 
     memset(L, 0, sizeof(*L));
 
@@ -71,14 +71,14 @@ static inline void htp_get_rows_vtcm_layout_build(
 
         L->src0_spad_half_size   = (src0_row_size + 255) & ~255;
         L->dst_spad_half_size    = (ne00 * sizeof(float) + 255) & ~255;
-        L->src0_bytes_per_thread = 2 * L->src0_spad_half_size;
-        L->dst_bytes_per_thread  = 2 * L->dst_spad_half_size;
+        L->src0_bytes_per_hvx    = 2 * L->src0_spad_half_size;
+        L->dst_bytes_per_hvx     = 2 * L->dst_spad_half_size;
     }
 
     size_t off = 0;
     VTCM_LAYOUT_ALLOC(off, off_src1, src1_alloc_bytes);
-    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
-    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread  * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_hvx * n_hvx);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_hvx  * n_hvx);
     L->total_bytes = off;
 }