Compare commits

..
Author SHA1 Message Date
geohot baa2f16bff revert 2025-10-09 18:31:06 +08:00
geohot ecf1477c0e merge those pattern matchers 2025-10-09 18:23:13 +08:00
geohot e4186add83 group div rules 2025-10-09 18:14:52 +08:00
385 changed files with 3882 additions and 52999 deletions
-14
View File
@@ -41,10 +41,6 @@ inputs:
description: "Install LLVM?"
required: false
default: 'false'
mesa:
description: "Install mesa"
required: false
default: 'false'
runs:
using: "composite"
steps:
@@ -293,13 +289,3 @@ runs:
if: inputs.llvm == 'true' && runner.os == 'macOS'
shell: bash
run: brew install llvm@20
# **** mesa ****
- name: Install mesa (linux)
if: inputs.mesa == 'true' && runner.os == 'Linux'
shell: bash
run: sudo curl -L https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
- name: Install mesa (macOS)
if: inputs.mesa == 'true' && runner.os == 'macOS'
shell: bash
run: brew install sirhcm/tinymesa/tinymesa
+1 -7
View File
@@ -36,9 +36,8 @@ jobs:
cuda: 'true'
webgpu: 'true'
llvm: 'true'
pydeps: 'pyyaml mako'
- name: Install autogen support packages
run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev llvm-20-dev
run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
- name: Verify OpenCL autogen
run: |
cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
@@ -90,8 +89,3 @@ jobs:
cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
./autogen_stubs.sh llvm
diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
- name: Verify mesa autogen
run: |
cp tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
./autogen_stubs.sh mesa
diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
+10 -11
View File
@@ -51,18 +51,17 @@ jobs:
rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
- name: reset process replay
run: python3.11 test/external/process_replay/reset.py
- name: Print macOS version
run: sw_vers
- name: Run Stable Diffusion
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
- name: Run Stable Diffusion without fp16
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=900 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
- name: Run Stable Diffusion v2
# TODO: very slow step time
run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=10000 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
# process replay can't capture this, the graph is too large
- name: Run SDXL
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
# TODO: too slow
# - name: Run SDXL
# run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
- name: Run model inference benchmark
run: METAL=1 python3.11 test/external/external_model_benchmark.py
- name: Test speed vs torch
@@ -626,14 +625,14 @@ jobs:
- name: benchmark openpilot 0.9.9 dmonitoring
run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
- name: openpilot compile3 0.9.9 driving_vision
run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=22 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot compile3 0.9.9 driving_policy
run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
- name: openpilot compile3 0.9.9 dmonitoring
run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=15 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
- name: openpilot compile3 Space Lab policy + vision
run: |
PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
- name: benchmark MobileNetV2 on DSP
run: |
@@ -642,7 +641,7 @@ jobs:
ln -s /data/home/tiny/tinygrad/testsig-*.so .
PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
# benchmark on DSP with NOOPT=1, the devectorizer has issues
PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
PYTHONPATH=. CC=clang-19 DSP=1 DONT_REALIZE_EXPAND=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
- name: Run process replay tests
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
- uses: actions/upload-artifact@v4
+1 -1
View File
@@ -12,7 +12,7 @@ jobs:
run_script_job:
runs-on: [self-hosted, Linux, tinybox]
if: github.repository_owner == 'tinygrad'
timeout-minutes: 720
timeout-minutes: 360
steps:
- name: Checkout Code
+73 -70
View File
@@ -89,65 +89,64 @@ jobs:
clang -O2 recognize.c -lm -o recognize
cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock
# TODO: fix the torch backend and reenable
# torchbackend:
# name: Torch Backend Tests
# runs-on: ubuntu-latest
# timeout-minutes: 15
# steps:
# - name: Checkout Code
# uses: actions/checkout@v4
# - name: Setup Environment
# uses: ./.github/actions/setup-tinygrad
# with:
# key: torch-backend-pillow-torchvision-et-pt
# deps: testing_minimal
# pydeps: "pillow torchvision expecttest"
# llvm: 'true'
# - name: Install ninja
# run: |
# sudo apt update || true
# sudo apt install -y --no-install-recommends ninja-build
# - name: Lint with ruff
# run: |
# pip3 install --upgrade --force-reinstall ruff==0.11.0
# python3 -m ruff check extra/torch_backend/backend.py
# - name: Test one op
# run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
# - name: Test ResNet-18
# run: DEBUG=2 python3 extra/torch_backend/example.py
# - name: My (custom) tests
# run: python3 extra/torch_backend/test.py
# - name: Test one op in torch tests
# run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
# - name: Test Ops with TINY_BACKEND
# run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
# - name: Test in-place operations on views
# run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
# - name: Test multi-gpu
# run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
torchbackend:
name: Torch Backend Tests
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: torch-backend-pillow-torchvision-et-pt
deps: testing_minimal
pydeps: "pillow torchvision expecttest"
llvm: 'true'
- name: Install ninja
run: |
sudo apt update || true
sudo apt install -y --no-install-recommends ninja-build
- name: Lint with ruff
run: |
pip3 install --upgrade --force-reinstall ruff==0.11.0
python3 -m ruff check extra/torch_backend/backend.py
- name: Test one op
run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
- name: Test ResNet-18
run: DEBUG=2 python3 extra/torch_backend/example.py
- name: My (custom) tests
run: python3 extra/torch_backend/test.py
- name: Test one op in torch tests
run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
- name: Test Ops with TINY_BACKEND
run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
- name: Test in-place operations on views
run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
- name: Test multi-gpu
run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
# torchbackendmore:
# name: Torch Backend Tests More
# runs-on: ubuntu-latest
# timeout-minutes: 15
# steps:
# - name: Checkout Code
# uses: actions/checkout@v4
# - name: Setup Environment
# uses: ./.github/actions/setup-tinygrad
# with:
# key: torch-backend-pillow-torchvision-et-pt
# deps: testing_minimal
# llvm: 'true'
# - name: Install ninja
# run: |
# sudo apt update || true
# sudo apt install -y --no-install-recommends ninja-build
# - name: Test beautiful_mnist in torch with TINY_BACKEND
# run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
# - name: Test some torch tests (expect failure)
# run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
torchbackendmore:
name: Torch Backend Tests More
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: torch-backend-pillow-torchvision-et-pt
deps: testing_minimal
llvm: 'true'
- name: Install ninja
run: |
sudo apt update || true
sudo apt install -y --no-install-recommends ninja-build
- name: Test beautiful_mnist in torch with TINY_BACKEND
run: CPU=1 CPU_LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
- name: Test some torch tests (expect failure)
run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
bepython:
name: Python Backend
@@ -239,6 +238,8 @@ jobs:
pip3 install --upgrade --force-reinstall ruff==0.11.0
python3 -m ruff check .
python3 -m ruff check examples/mlperf/ --ignore E501
- name: Lint tinygrad with pylint
run: python -m pylint tinygrad/
- name: Run mypy
run: |
python -m mypy --strict-equality --lineprecision-report .
@@ -273,8 +274,6 @@ jobs:
# run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
- name: Run Clip tests for SD MLPerf on NULL backend
run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
- name: Run AMD emulated BERT training on NULL backend
run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
# TODO: support fake weights
#- name: Run LLaMA 7B on 4 fake devices
# run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -310,6 +309,10 @@ jobs:
run: python test/external/fuzz_symbolic.py
- name: Fuzz Test fast idiv
run: python test/external/fuzz_fast_idiv.py
- name: Fuzz Test shapetracker
run: |
python test/external/fuzz_shapetracker.py
python test/external/fuzz_shapetracker_math.py
- name: Fuzz Test shape ops
run: python test/external/fuzz_shape_ops.py
@@ -374,7 +377,7 @@ jobs:
llvm: 'true'
- name: Test openpilot model kernel count and gate usage
run: |
ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2081 ALLOWED_GATED_READ_IMAGE=28 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=543 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
- name: Test openpilot alt model correctness (float32)
run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
- name: Test openpilot fastvits model correctness (float32)
@@ -383,6 +386,8 @@ jobs:
# run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
- name: Test openpilot LLVM compile
run: CPU=1 CPU_LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
- name: Test openpilot compile4
run: NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -672,7 +677,7 @@ jobs:
strategy:
fail-fast: false
matrix:
backend: [llvm, cpu, opencl, lvp]
backend: [llvm, cpu, opencl]
name: Linux (${{ matrix.backend }})
runs-on: ubuntu-22.04
@@ -686,10 +691,9 @@ jobs:
key: ${{ matrix.backend }}-minimal
deps: testing_minimal
opencl: ${{ matrix.backend == 'opencl' && 'true' }}
llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
mesa: ${{ matrix.backend == 'lvp' && 'true' }}
llvm: ${{ matrix.backend == 'llvm' && 'true' }}
- name: Set env
run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' }}" >> $GITHUB_ENV
- name: Check Device.DEFAULT and print some source
run: |
python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
@@ -891,7 +895,7 @@ jobs:
strategy:
fail-fast: false
matrix:
backend: [metal, llvm, cpu, lvp]
backend: [metal, llvm, cpu]
name: MacOS (${{ matrix.backend }})
runs-on: macos-15
timeout-minutes: 20
@@ -904,13 +908,12 @@ jobs:
key: macos-${{ matrix.backend }}-minimal
deps: testing_minimal
pydeps: "capstone"
llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
mesa: ${{ matrix.backend == 'lvp' && 'true' }}
llvm: ${{ matrix.backend == 'llvm' && 'true' }}
- name: Set env
run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
- name: Check Device.DEFAULT and print some source
run: |
python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
- name: Run pytest (${{ matrix.backend }})
run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
-1
View File
@@ -38,7 +38,6 @@ extra/huggingface_onnx/models/*
extra/huggingface_onnx/*.yaml
extra/weights
venv
venv_sd_mlperf
examples/**/net.*[js,json]
examples/**/*.safetensors
node_modules
+10 -4
View File
@@ -20,15 +20,21 @@ repos:
language: system
always_run: true
pass_filenames: false
- id: tests
name: subset of tests
entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
language: system
always_run: true
pass_filenames: false
- id: example
name: test all devices
name: multi device tests
entry: python3 test/external/external_test_example.py
language: system
always_run: true
pass_filenames: false
- id: tests
name: subset of tests
entry: env PYTHONPATH="." python3 -m pytest -n=8 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
- id: pylint
name: pylint
entry: python3 -m pylint tinygrad/
language: system
always_run: true
pass_filenames: false
+4
View File
@@ -30,6 +30,10 @@ persistent=yes
# Specify a configuration file.
#rcfile=
# When enabled, pylint would attempt to guess common misconfiguration and emit
# user-friendly hints instead of false-positive error messages
suggestion-mode=yes
# Allow loading of arbitrary C extensions. Extensions are imported into the
# active Python interpreter and may run arbitrary code.
unsafe-load-any-extension=no
+2 -101
View File
@@ -414,29 +414,10 @@ generate_sqtt() {
clang2py -k cdefstum \
extra/sqtt/sqtt.h \
-o $BASE/sqtt.py
fixup $BASE/sqtt.py
sed -i "s\import ctypes\import ctypes, os\g" $BASE/sqtt.py
python3 -c "import tinygrad.runtime.autogen.sqtt"
ROCPROF_COMMIT_HASH=dd0485100971522cc4cd8ae136bdda431061a04d
ROCPROF_SRC=/tmp/rocprof-trace-decoder-$ROCPROF_COMMIT_HASH
if [ ! -d "$ROCPROF_SRC" ]; then
git clone https://github.com/ROCm/rocprof-trace-decoder $ROCPROF_SRC
pushd .
cd $ROCPROF_SRC
git reset --hard $ROCPROF_COMMIT_HASH
popd
fi
clang2py -k cdefstum \
$ROCPROF_SRC/include/rocprof_trace_decoder.h \
$ROCPROF_SRC/include/trace_decoder_instrument.h \
$ROCPROF_SRC/include/trace_decoder_types.h \
-o extra/sqtt/rocprof/rocprof.py
fixup extra/sqtt/rocprof/rocprof.py
sed -i '1s/^/# pylint: skip-file\n/' extra/sqtt/rocprof/rocprof.py
sed -i "s/import ctypes/import ctypes, ctypes.util/g" extra/sqtt/rocprof/rocprof.py
sed -i "s|FunctionFactoryStub()|ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder'))|g" extra/sqtt/rocprof/rocprof.py
}
generate_webgpu() {
@@ -461,85 +442,6 @@ generate_libusb() {
python3 -c "import tinygrad.runtime.autogen.libusb"
}
generate_mesa() {
MESA_TAG="mesa-25.2.4"
MESA_SRC=/tmp/mesa-$MESA_TAG
TINYMESA_TAG=tinymesa-32dc66c
TINYMESA_DIR=/tmp/tinymesa-$MESA_TAG-$TINYMESA_TAG/
TINYMESA_SO=$TINYMESA_DIR/libtinymesa_cpu.so
if [ ! -d "$MESA_SRC" ]; then
git clone --depth 1 --branch $MESA_TAG https://gitlab.freedesktop.org/mesa/mesa.git $MESA_SRC
pushd .
cd $MESA_SRC
git reset --hard $MESA_COMMIT_HASH
# clang 14 doesn't support packed enums
sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/headers/nv_device_info.h
sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/compiler/nak.h
sed -i "s/nir_instr_type \(\w\+\);/uint8_t \1;/" $MESA_SRC/src/compiler/nir/nir.h
mkdir -p gen/util/format
python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h
python3 src/compiler/nir/nir_opcodes_h.py > gen/nir_opcodes.h
python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
python3 src/compiler/nir/nir_intrinsics_indices_h.py --outdir gen
python3 src/compiler/nir/nir_builder_opcodes_h.py > gen/nir_builder_opcodes.h
python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
python3 src/compiler/builtin_types_h.py gen/builtin_types.h
popd
fi
if [ ! -d "$TINYMESA_DIR" ]; then
mkdir $TINYMESA_DIR
curl -L https://github.com/sirhcm/tinymesa/releases/download/$TINYMESA_TAG/libtinymesa_cpu-$MESA_TAG-linux-amd64.so -o $TINYMESA_SO
fi
clang2py -k cdefstu \
$MESA_SRC/src/compiler/nir/nir.h \
$MESA_SRC/src/compiler/nir/nir_builder.h \
$MESA_SRC/src/compiler/nir/nir_shader_compiler_options.h \
$MESA_SRC/src/compiler/nir/nir_serialize.h \
$MESA_SRC/gen/nir_intrinsics.h \
$MESA_SRC/src/nouveau/headers/nv_device_info.h \
$MESA_SRC/src/nouveau/compiler/nak.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_passmgr.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_misc.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_type.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_init.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_nir.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_struct.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_jit_types.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_flow.h \
$MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_const.h \
$MESA_SRC/src/compiler/glsl_types.h \
$MESA_SRC/src/util/blob.h \
$MESA_SRC/src/util/ralloc.h \
--clang-args="-DHAVE_ENDIAN_H -DHAVE_STRUCT_TIMESPEC -DHAVE_PTHREAD -I$MESA_SRC/src -I$MESA_SRC/include -I$MESA_SRC/gen -I$MESA_SRC/src/compiler/nir -I$MESA_SRC/src/gallium/auxiliary -I$MESA_SRC/src/gallium/include -I$(llvm-config-20 --includedir)" \
-l $TINYMESA_SO \
-o $BASE/mesa.py
LVP_NIR_OPTIONS=$(./extra/mesa/lvp_nir_options.sh $MESA_SRC)
fixup $BASE/mesa.py
patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "f'{brew_prefix()}/lib/libtinymesa_cpu.dylib'"
echo "lvp_nir_options = gzip.decompress(base64.b64decode('$LVP_NIR_OPTIONS'))" >> $BASE/mesa.py
cat <<EOF | sed -i "/import ctypes.*/r /dev/stdin" $BASE/mesa.py
def brew_prefix():
try: return subprocess.check_output(['brew', '--prefix', 'tinymesa']).decode().strip()
except Exception: return ''
EOF
sed -i "/in_dll/s/.*/try: &\nexcept AttributeError: pass/" $BASE/mesa.py
sed -i "s/import ctypes/import ctypes, ctypes.util, os, gzip, base64, subprocess, tinygrad.helpers as helpers/" $BASE/mesa.py
sed -i "s/ctypes.CDLL('.\+')/(dll := _try_dlopen_tinymesa_cpu())/" $BASE/mesa.py
echo "def __getattr__(nm): raise AttributeError() if dll else FileNotFoundError(f'libtinymesa not found (MESA_PATH={BASE}). See https://github.com/sirhcm/tinymesa ($TINYMESA_TAG, $MESA_TAG)')" >> $BASE/mesa.py
sed -i "s/ctypes.glsl_base_type/glsl_base_type/" $BASE/mesa.py
# bitfield bug in clang2py
sed -i "s/('fp_fast_math', ctypes.c_bool, 9)/('fp_fast_math', ctypes.c_uint32, 9)/" $BASE/mesa.py
sed -i "s/('\(\w\+\)', pipe_shader_type, 8)/('\1', ctypes.c_ubyte)/" $BASE/mesa.py
sed -i "s/\([0-9]\+\)()/\1/" $BASE/mesa.py
sed -i "s/\(struct_nir_builder._pack_\) = 1/\1 = 0/" $BASE/mesa.py
python3 -c "import tinygrad.runtime.autogen.mesa"
}
if [ "$1" == "opencl" ]; then generate_opencl
elif [ "$1" == "hip" ]; then generate_hip
elif [ "$1" == "comgr" ]; then generate_comgr
@@ -563,7 +465,6 @@ elif [ "$1" == "pci" ]; then generate_pci
elif [ "$1" == "vfio" ]; then generate_vfio
elif [ "$1" == "webgpu" ]; then generate_webgpu
elif [ "$1" == "libusb" ]; then generate_libusb
elif [ "$1" == "mesa" ]; then generate_mesa
elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu; generate_mesa
elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu
else echo "usage: $0 <type>"
fi
+1 -1
View File
@@ -232,7 +232,7 @@ if __name__ == "__main__":
gpt2 = GPT2.build_gguf(args.model_size) if args.model_size.startswith("gpt2_gguf_") else GPT2.build(args.model_size)
if args.benchmark != -1:
gpt2.model(Tensor.randint(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
gpt2.model(Tensor.rand(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
else:
texts = gpt2.generate(args.prompt, args.count, args.temperature, timing=args.timing, batch_size=args.batch_size)
if not args.noshow:
+2 -4
View File
@@ -1188,9 +1188,7 @@ def train_bert():
if MLLOGGER and RUNMLPERF:
MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i})
if getenv("RESET_STEP"): train_step_bert.reset()
elif getenv("FREE_INTERMEDIATE", 0) and train_step_bert.captured is not None:
# TODO: FREE_INTERMEDIATE nan'ed after jit step 2
train_step_bert.captured.free_intermediates()
elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
eval_lm_losses = []
eval_clsf_losses = []
eval_lm_accs = []
@@ -1224,7 +1222,7 @@ def train_bert():
return
if getenv("RESET_STEP"): eval_step_bert.reset()
elif getenv("FREE_INTERMEDIATE", 0) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
elif getenv("FREE_INTERMEDIATE", 1) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
del eval_data
avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
@@ -1,17 +0,0 @@
#!/bin/bash
export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
export IGNORE_OOB=1
export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
# export BEAM_LOG_SURPASS_MAX=1
# export BASEDIR="/raid/datasets/wiki"
export RESET_STEP=1
export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
python3 examples/mlperf/model_train.py
@@ -1,69 +0,0 @@
# 1. Problem
This problem uses BERT for NLP.
## Requirements
Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
```
git clone https://github.com/tinygrad/tinygrad.git
python3 -m pip install -e ".[mlperf]"
```
Also install gdown (for dataset), numpy, tqdm and tensorflow.
```
pip install gdown numpy tqdm tensorflow
```
### tinybox_green
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
This is the default on production tinybox green.
# 2. Directions
## Steps to download and verify data
### 1. Download raw data
```
BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
```
### 2. Preprocess train and validation data
Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended.
#### Training:
```
BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
```
Generating a specific topic (Between 0 and 499)
```
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
```
#### Validation:
```
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
```
## Running
### tinybox_green
#### Steps to run benchmark
```
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
```
### tinybox_red
#### Steps to run benchmark
```
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
```
### tinybox_8xMI300X
#### Steps to run benchmark
```
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
```
@@ -1,17 +0,0 @@
#!/bin/bash
export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
export BASEDIR="/raid/datasets/wiki"
export BENCHMARK=10 BERT_LAYERS=2
python3 examples/mlperf/model_train.py
@@ -1,20 +0,0 @@
#!/bin/bash
export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
export BASEDIR="/raid/datasets/wiki"
export WANDB=1 PARALLEL=0
RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -1,31 +0,0 @@
#!/bin/bash
set -e # Exit on any error
set -o pipefail # Make pipeline fail if any command fails
export PYTHONPATH="." AMD=1
export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_8xMI300X"
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
export BASEDIR="/raid/datasets/wiki"
# pip install -e ".[mlperf]"
export LOGMLPERF=1
export SEED=$RANDOM
DATETIME=$(date "+%m%d%H%M")
LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log"
BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
# run
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,69 +0,0 @@
# 1. Problem
This problem uses BERT for NLP.
## Requirements
Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
```
git clone https://github.com/tinygrad/tinygrad.git
python3 -m pip install -e ".[mlperf]"
```
Also install gdown (for dataset), numpy, tqdm and tensorflow.
```
pip install gdown numpy tqdm tensorflow
```
### tinybox_green
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
This is the default on production tinybox green.
# 2. Directions
## Steps to download and verify data
### 1. Download raw data
```
BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
```
### 2. Preprocess train and validation data
Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended.
#### Training:
```
BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
```
Generating a specific topic (Between 0 and 499)
```
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
```
#### Validation:
```
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
```
## Running
### tinybox_green
#### Steps to run benchmark
```
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
```
### tinybox_red
#### Steps to run benchmark
```
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
```
### tinybox_8xMI300X
#### Steps to run benchmark
```
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
```
@@ -1,17 +0,0 @@
#!/bin/bash
export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
export BEAM_LOG_SURPASS_MAX=1
export BASEDIR="/raid/datasets/wiki"
export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
python3 examples/mlperf/model_train.py
@@ -1,16 +0,0 @@
#!/bin/bash
export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
export BASEDIR="/raid/datasets/wiki"
export WANDB=1 PARALLEL=0
RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -1,28 +0,0 @@
#!/bin/bash
set -e # Exit on any error
set -o pipefail # Make pipeline fail if any command fails
export PYTHONPATH="." NV=1
export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_green"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
export BASEDIR="/raid/datasets/wiki"
# pip install -e ".[mlperf]"
export LOGMLPERF=1
export SEED=$RANDOM
DATETIME=$(date "+%m%d%H%M")
LOGFILE="bert_green_${DATETIME}_${SEED}.log"
# init
BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
# run
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,69 +0,0 @@
# 1. Problem
This problem uses BERT for NLP.
## Requirements
Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
```
git clone https://github.com/tinygrad/tinygrad.git
python3 -m pip install -e ".[mlperf]"
```
Also install gdown (for dataset), numpy, tqdm and tensorflow.
```
pip install gdown numpy tqdm tensorflow
```
### tinybox_green
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
This is the default on production tinybox green.
# 2. Directions
## Steps to download and verify data
### 1. Download raw data
```
BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
```
### 2. Preprocess train and validation data
Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended.
#### Training:
```
BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
```
Generating a specific topic (Between 0 and 499)
```
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
```
#### Validation:
```
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
```
## Running
### tinybox_green
#### Steps to run benchmark
```
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
```
### tinybox_red
#### Steps to run benchmark
```
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
```
### tinybox_8xMI300X
#### Steps to run benchmark
```
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
```
@@ -1,18 +0,0 @@
#!/bin/bash
export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
export BEAM_LOG_SURPASS_MAX=1
export BASEDIR="/raid/datasets/wiki"
export RESET_STEP=1
export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
python3 examples/mlperf/model_train.py
@@ -1,16 +0,0 @@
#!/bin/bash
export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
export BASEDIR="/raid/datasets/wiki"
export WANDB=1 PARALLEL=0
RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -1,31 +0,0 @@
#!/bin/bash
set -e # Exit on any error
set -o pipefail # Make pipeline fail if any command fails
export PYTHONPATH="." AMD=1
export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_red"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
export BASEDIR="/raid/datasets/wiki"
# pip install -e ".[mlperf]"
export LOGMLPERF=1
export SEED=$RANDOM
DATETIME=$(date "+%m%d%H%M")
LOGFILE="bert_red_${DATETIME}_${SEED}.log"
export HCQDEV_WAIT_TIMEOUT_MS=100000 # prevents hang?
# init
sleep 5 && sudo rmmod amdgpu || true
BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
# run
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,50 +0,0 @@
# 1. Problem
This problem uses the ResNet-50 CNN to do image classification.
## Requirements
Install tinygrad and mlperf-logging from master.
```
git clone https://github.com/tinygrad/tinygrad.git
python3 -m pip install -e ".[mlperf]"
```
### tinybox_green
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
This is the default on production tinybox green.
### tinybox_red
Disable cwsr
This is the default on production tinybox red.
```
sudo vi /etc/modprobe.d/amdgpu.conf
cat <<EOF > /etc/modprobe.d/amdgpu.conf
options amdgpu cwsr_enable=0
EOF
sudo update-initramfs -u
sudo reboot
# validate
sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
```
# 2. Directions
## Steps to download and verify data
```
IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
```
## Steps for one time setup
### tinybox_red
```
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
```
## Steps to run benchmark
```
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
```
@@ -1,13 +0,0 @@
#!/bin/bash
export PYTHONPATH="." NV=1
export MODEL="resnet"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
export RESET_STEP=0
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
export BENCHMARK=10 DEBUG=2
python3 examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
#!/bin/bash
export PYTHONPATH="." NV=1
export MODEL="resnet"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
export RESET_STEP=0
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
export EVAL_START_EPOCH=3 EVAL_FREQ=4
export WANDB=1 PARALLEL=0
python3 examples/mlperf/model_train.py
@@ -1,25 +0,0 @@
#!/bin/bash
set -e # Exit on any error
set -o pipefail # Make pipeline fail if any command fails
export PYTHONPATH="." NV=1
export MODEL="resnet"
export SUBMISSION_PLATFORM="tinybox_green"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
export RESET_STEP=0
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
# pip install -e ".[mlperf]"
export LOGMLPERF=${LOGMLPERF:-1}
export SEED=$RANDOM
DATETIME=$(date "+%m%d%H%M")
LOGFILE="resnet_green_${DATETIME}_${SEED}.log"
# init
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
# run
PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,50 +0,0 @@
# 1. Problem
This problem uses the ResNet-50 CNN to do image classification.
## Requirements
Install tinygrad and mlperf-logging from master.
```
git clone https://github.com/tinygrad/tinygrad.git
python3 -m pip install -e ".[mlperf]"
```
### tinybox_green
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
This is the default on production tinybox green.
### tinybox_red
Disable cwsr
This is the default on production tinybox red.
```
sudo vi /etc/modprobe.d/amdgpu.conf
cat <<EOF > /etc/modprobe.d/amdgpu.conf
options amdgpu cwsr_enable=0
EOF
sudo update-initramfs -u
sudo reboot
# validate
sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
```
# 2. Directions
## Steps to download and verify data
```
IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
```
## Steps for one time setup
### tinybox_red
```
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
```
## Steps to run benchmark
```
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
```
@@ -1,13 +0,0 @@
#!/bin/bash
export PYTHONPATH="." AMD=1
export MODEL="resnet"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
export RESET_STEP=0
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
export BENCHMARK=10 DEBUG=${DEBUG:-2}
python3 examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
#!/bin/bash
export PYTHONPATH="." AMD=1
export MODEL="resnet"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
export RESET_STEP=0
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
export EVAL_START_EPOCH=3 EVAL_FREQ=4
export WANDB=1 PARALLEL=0
python3 examples/mlperf/model_train.py
@@ -1,26 +0,0 @@
#!/bin/bash
set -e # Exit on any error
set -o pipefail # Make pipeline fail if any command fails
export PYTHONPATH="." AMD=1
export MODEL="resnet"
export SUBMISSION_PLATFORM="tinybox_red"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
export RESET_STEP=0
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
# pip install -e ".[mlperf]"
export LOGMLPERF=${LOGMLPERF:-1}
export SEED=$RANDOM
DATETIME=$(date "+%m%d%H%M")
LOGFILE="resnet_red_${DATETIME}_${SEED}.log"
# init
sleep 5 && sudo rmmod amdgpu || true
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
# run
PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,8 +0,0 @@
#!/bin/bash
rocm-smi --setprofile compute
rocm-smi --setmclk 3
rocm-smi --setperflevel high
# power cap to 350W
echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
@@ -1,38 +0,0 @@
# 1. Problem
This problem uses RetinaNet for SSD.
## Requirements
Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
```
git clone https://github.com/tinygrad/tinygrad.git
python3 -m pip install -e ".[mlperf]"
```
Also install the following dependencies:
```
pip install tqdm numpy pycocotools boto3 pandas torch torchvision
```
### tinybox_green
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
This is the default on production tinybox green.
# 2. Directions
## Steps to download data
Run the following:
```
BASEDIR=/raid/datasets/openimages python3 extra/datasets/openimages.py
```
## Running
### tinybox_green
#### Steps to run benchmark
```
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh
```
@@ -1,14 +0,0 @@
#!/bin/bash
export PYTHONPATH="." NV=1
export MODEL="retinanet"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
export BASEDIR="/raid/datasets/openimages"
# export RESET_STEP=0
export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
export BENCHMARK=5 DEBUG=2
python examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
#!/bin/bash
export PYTHONPATH="." NV=1
export MODEL="retinanet"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
export BASEDIR="/raid/datasets/openimages"
# export RESET_STEP=0
export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
export WANDB=1 PARALLEL=0
export RUNMLPERF=1
python examples/mlperf/model_train.py
@@ -1,25 +0,0 @@
#!/bin/bash
set -e # Exit on any error
set -o pipefail # Make pipeline fail if any command fails
export PYTHONPATH="." NV=1
export MODEL="retinanet"
export SUBMISSION_PLATFORM="tinybox_green"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
export TRAIN_BEAM=2 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
export IGNORE_JIT_FIRST_BEAM=1
export BASEDIR="/raid/datasets/openimages"
# pip install -e ".[mlperf]"
export LOGMLPERF=1
export SEED=$RANDOM
DATETIME=$(date "+%m%d%H%M")
LOGFILE="retinanet_green_${DATETIME}_${SEED}.log"
# init
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
# run
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,14 +0,0 @@
#!/bin/bash
export PYTHONPATH="." AMD=1
export MODEL="retinanet"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
export BASEDIR="/raid/datasets/openimages"
# export RESET_STEP=0
export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
export BENCHMARK=5 DEBUG=2
python examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
#!/bin/bash
export PYTHONPATH="." AMD=1
export MODEL="retinanet"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
export BASEDIR="/raid/datasets/openimages"
# export RESET_STEP=0
export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
export WANDB=1 PARALLEL=0
export RUNMLPERF=1
python examples/mlperf/model_train.py
@@ -1,38 +0,0 @@
{
"submitter": "tinycorp",
"division": "closed",
"status": "Available on-premise",
"system_name": "tinybox 8xMI300X",
"number_of_nodes": "1",
"host_processors_per_node": "2",
"host_processor_model_name": "AMD EPYC 9354",
"host_processor_core_count": "32",
"host_processor_vcpu_count": "64",
"host_processor_frequency": "",
"host_processor_caches": "",
"host_processor_interconnect": "",
"host_memory_capacity": "2304GB",
"host_storage_type": "NVMe SSD",
"host_storage_capacity": "3x 4TB raid array",
"host_networking": "",
"host_networking_topology": "",
"host_memory_configuration": "24x 96GB DDR5",
"accelerators_per_node": "8",
"accelerator_model_name": "AMD Instinct MI300X 192GB HBM3",
"accelerator_host_interconnect": "PCIe 5.0 x16",
"accelerator_frequency": "",
"accelerator_on-chip_memories": "",
"accelerator_memory_configuration": "HBM3",
"accelerator_memory_capacity": "192GB",
"accelerator_interconnect": "",
"accelerator_interconnect_topology": "",
"cooling": "air",
"hw_notes": "",
"framework": "tinygrad, branch mlperf_training_v5.0",
"other_software_stack": {
"python": "3.10.16",
"ROCm": "3.0.0+94441cb"
},
"operating_system": "Ubuntu 24.04.1 LTS",
"sw_notes": ""
}
@@ -1,38 +0,0 @@
{
"submitter": "tinycorp",
"division": "closed",
"status": "Available on-premise",
"system_name": "tinybox green",
"number_of_nodes": "1",
"host_processors_per_node": "1",
"host_processor_model_name": "AMD EPYC 7532",
"host_processor_core_count": "32",
"host_processor_vcpu_count": "64",
"host_processor_frequency": "",
"host_processor_caches": "",
"host_processor_interconnect": "",
"host_memory_capacity": "128GB",
"host_storage_type": "NVMe SSD",
"host_storage_capacity": "4 TB raid array + 1 TB boot",
"host_networking": "",
"host_networking_topology": "",
"host_memory_configuration": "8x 16GB DDR4",
"accelerators_per_node": "6",
"accelerator_model_name": "NVIDIA GeForce RTX 4090",
"accelerator_host_interconnect": "PCIe 4.0 x16",
"accelerator_frequency": "",
"accelerator_on-chip_memories": "",
"accelerator_memory_configuration": "GDDR6X",
"accelerator_memory_capacity": "24GB",
"accelerator_interconnect": "",
"accelerator_interconnect_topology": "",
"cooling": "air",
"hw_notes": "",
"framework": "tinygrad, branch mlperf_training_v5.0",
"other_software_stack": {
"python": "3.10.12",
"CUDA": "12.4"
},
"operating_system": "Ubuntu 22.04.4",
"sw_notes": ""
}
@@ -1,37 +0,0 @@
{
"submitter": "tinycorp",
"division": "closed",
"status": "Available on-premise",
"system_name": "tinybox red",
"number_of_nodes": "1",
"host_processors_per_node": "1",
"host_processor_model_name": "AMD EPYC 7532",
"host_processor_core_count": "32",
"host_processor_vcpu_count": "64",
"host_processor_frequency": "",
"host_processor_caches": "",
"host_processor_interconnect": "",
"host_memory_capacity": "128GB",
"host_storage_type": "NVMe SSD",
"host_storage_capacity": "4 TB raid array + 1 TB boot",
"host_networking": "",
"host_networking_topology": "",
"host_memory_configuration": "8x 16GB DDR4",
"accelerators_per_node": "6",
"accelerator_model_name": "AMD Radeon RX 7900 XTX",
"accelerator_host_interconnect": "PCIe 4.0 x16",
"accelerator_frequency": "",
"accelerator_on-chip_memories": "",
"accelerator_memory_configuration": "GDDR6",
"accelerator_memory_capacity": "24GB",
"accelerator_interconnect": "",
"accelerator_interconnect_topology": "",
"cooling": "air",
"hw_notes": "",
"framework": "tinygrad, branch mlperf_training_v5.0",
"other_software_stack": {
"python": "3.10.12"
},
"operating_system": "Ubuntu 22.04.4",
"sw_notes": ""
}
+47
View File
@@ -0,0 +1,47 @@
import sys
from tinygrad import Tensor, fetch, GlobalCounters, dtypes
from tinygrad.uop.ops import UOp
from tinygrad.nn.onnx import OnnxRunner
from tinygrad.schedule.rangeify import get_rangeify_map
from tinygrad.engine.schedule import create_schedule_with_vars
from tinygrad.engine.realize import run_schedule
# NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
if __name__ == "__main__":
onnx_file = fetch(OPENPILOT_MODEL)
run_onnx = OnnxRunner(onnx_file)
inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
root = out.uop
targets = [x.uop for x in inputs.values()]
print(targets)
# TODO: abstract this from gradient?
# compute the target path (top down)
in_target_path: dict[UOp, bool] = {}
for u in root.toposort(): in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
independent_set = {}
for u in root.toposort():
if in_target_path[u]:
for s in u.src:
if not in_target_path[s]:
independent_set[s] = None
independent = UOp.sink(*independent_set.keys())
kernelized = get_rangeify_map(independent)
independent = independent.substitute(kernelized)
schedule, var_vals = create_schedule_with_vars(independent)
run_schedule(schedule)
print("**** real ****")
GlobalCounters.reset()
out.uop = root.substitute(kernelized)
out.kernelize()
# realize
out.realize()
+1 -2
View File
@@ -269,8 +269,7 @@ if __name__ == "__main__":
# load in weights
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
model_bin = fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt')
load_state_dict(model, torch_load(model_bin)['state_dict'], verbose=False, strict=False, realize=False)
load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], verbose=False, strict=False, realize=False)
if args.fp16:
for k,v in get_state_dict(model).items():
+2 -2
View File
@@ -19,8 +19,8 @@ from tinygrad.helpers import fetch, getenv
# QUANT=1 python3 examples/test_onnx_imagenet.py
# https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx
# python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
# VIZ=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx
# DONT_REALIZE_EXPAND=1 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
# VIZ=1 DONT_REALIZE_EXPAND=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx
def imagenet_dataloader(cnt=0):
input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
-23
View File
@@ -1,23 +0,0 @@
#!/bin/sh
if [ "$#" -ne 1 ] || ! [ -d $1 ]; then
echo "usage: $0 MESA_PREFIX"
exit 1
fi
TMP=$(mktemp)
trap 'rm -f "$TMP"' EXIT
(
cat <<EOF
#define HAVE_ENDIAN_H
#define HAVE_STRUCT_TIMESPEC
#define HAVE_PTHREAD
#include <unistd.h>
#include "nir_shader_compiler_options.h"
#include "compiler/shader_enums.h"
EOF
sed -n '/struct nir_shader_compiler_options/,/^}/{p;/^}/q}' $1/src/gallium/drivers/llvmpipe/lp_screen.c
echo "int main(void) { write(1, &gallivm_nir_options, sizeof(gallivm_nir_options)); }"
) | cc -x c -o $TMP - -I$1/src/compiler/nir -I$1/src -I$1/include && $TMP | gzip | base64 -w0
+4 -107
View File
@@ -3,21 +3,8 @@ from tinygrad.tensor import _to_np_dtype
from tinygrad.nn.onnx import OnnxRunner, OnnxValue
import numpy as np
import onnxruntime as ort
ort_options = ort.SessionOptions()
ort_options.log_severity_level = 3
def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
"""
Generate example input tensors based on the provided ONNX graph input specifications.
NOTE: This is not guaranteed to be reliable. It's a best-effort helper
that uses heuristics to guess input shapes and values.
Example:
from tinygrad.nn.onnx import OnnxRunner
from extra.onnx_helpers import get_example_inputs
inputs = get_example_inputs(OnnxRunner(model_path).graph_inputs)
"""
def _get_shape(onnx_shape: tuple[str|int]):
shape = []
for onnx_dim in onnx_shape:
@@ -57,9 +44,11 @@ def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
ret.update({name:value})
return ret
def _get_tinygrad_and_ort_np_outputs(onnx_file, inputs):
def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
run_onnx = OnnxRunner(onnx_file)
ort_options = ort.SessionOptions()
ort_options.log_severity_level = 3
ort_sess = ort.InferenceSession(onnx_file, ort_options, ["CPUExecutionProvider"])
np_inputs = {k:v.numpy() if isinstance(v, Tensor) else v for k,v in inputs.items()}
out_names = list(run_onnx.graph_outputs)
@@ -67,101 +56,9 @@ def _get_tinygrad_and_ort_np_outputs(onnx_file, inputs):
ort_out = dict(zip(out_names, out_values))
tinygrad_out = run_onnx(inputs)
Tensor.realize(*(x for x in tinygrad_out.values() if x is not None))
tinygrad_out = {k:v.numpy() if v is not None else None for k,v in tinygrad_out.items()}
return tinygrad_out, ort_out
def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
"""
Compares the final output tensors of an onnx model run in tinygrad and onnxruntime.
"""
tinygrad_out, ort_out = _get_tinygrad_and_ort_np_outputs(onnx_file, inputs)
assert tinygrad_out.keys() == ort_out.keys()
for k in tinygrad_out.keys():
tiny_v, onnx_v = tinygrad_out[k], ort_out[k]
if tiny_v is None: assert onnx_v is None, f"{k}: {tiny_v=}, {onnx_v=}"
else: np.testing.assert_allclose(tiny_v, onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tinygrad_out.keys()}")
def validate_all_intermediates(onnx_file, inputs, rtol=1e-5, atol=1e-5):
"""
Compares all intermediate node output of an onnx model run in tinygrad and onnxruntime.
"""
report = generate_node_output_report(onnx_file, inputs)
for i, node in enumerate(report):
node_name = node["node"]
op = node["op"]
outputs = node["outputs"]
for output in outputs:
output_name = output["name"]
tinygrad_out = output["tinygrad"]
ort_out = output["onnxruntime"]
try:
if tinygrad_out is None: assert ort_out is None, f"None outputs are not equal {tinygrad_out=} {ort_out=}"
else: np.testing.assert_allclose(tinygrad_out, ort_out, rtol=rtol, atol=atol)
print(f"Validated {i}: {op=} {node_name=} {output_name=}")
except AssertionError as e:
print(f"FAILED {i}: {op=} {node_name=} {output_name=}")
print(str(e).strip() + "\n")
def generate_node_output_report(onnx_file, inputs):
"""
Build a report of all ONNX node outputs from tinygrad and onnxruntime
Returns:
A list of dictionaries, where each entry corresponds to one
node in the ONNX graph. The structure is as follows:
[
{
"node": str, # The name of the ONNX node.
"op": str, # The operation type of the ONNX node.
"outputs": [
{
"name": str, # The name of the output tensor.
"tinygrad": np.ndarray | None, # The output value from tinygrad.
"onnxruntime": np.ndarray | None, # The output value from onnxruntime.
},
...
]
},
...
]
"""
import onnx_graphsurgeon as gs
import onnx
import tempfile
# rewrite the model to output all the node outputs
# `infer_shapes` here tries to fill the shapes and dtypes of intermediate values which graphsurgeon requires when assigning them as outputs
inferred_model = onnx.shape_inference.infer_shapes(onnx.load(onnx_file))
model = gs.import_onnx(inferred_model)
model_nodes = model.nodes
node_outputs = [n.outputs for n in model.nodes]
model.outputs = [
each_output for outputs in node_outputs for each_output in outputs
if not (each_output.dtype is None and each_output.shape is None) # output with None dtype and None shape is likely a `None` value
]
rewritten_model = gs.export_onnx(model)
# TODO: remove this once ORT supports 1.18.0
if getattr(rewritten_model, "ir_version", 0) > 10:
rewritten_model.ir_version = 10
with tempfile.NamedTemporaryFile(suffix=".onnx") as f:
onnx.save(rewritten_model, f.name)
rewritten_model_path = f.name
tinygrad_out, ort_out = _get_tinygrad_and_ort_np_outputs(rewritten_model_path, inputs)
report = []
for node in model_nodes:
outputs = []
for each_output in node.outputs:
if each_output.dtype is None and each_output.shape is None:
continue
name = each_output.name
tinygrad_output = tinygrad_out[name]
ort_output = ort_out[name]
outputs.append({"name": name, "tinygrad": tinygrad_output, "onnxruntime": ort_output})
report.append({"node": node.name, "op": node.op, "outputs": outputs})
return report
else: np.testing.assert_allclose(tiny_v.numpy(), onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tinygrad_out.keys()}")
+1 -1
View File
@@ -81,7 +81,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
ret = [float(x) for x in ret]
if use_sts:
my_sts = dedup([(x.shape == lin.full_shape, x.is_expanded(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
my_sts = dedup([(x.shape == lin.full_shape, x.real_strides(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
assert len(my_sts) < MAX_BUFS
sts_len = 3 + 5*MAX_DIMS
for s in my_sts:
-68
View File
@@ -1,68 +0,0 @@
import ctypes
from dataclasses import dataclass
import tinygrad.runtime.autogen.comgr as comgr
from tinygrad.runtime.support.compiler_amd import check
@dataclass
class InstrCtx:
pc:int=0
inst:str=""
@comgr.amd_comgr_create_disassembly_info.argtypes[2]
def instr_cb(text, user_data):
c = ctypes.cast(user_data, ctypes.POINTER(ctypes.py_object)).contents.value
c.inst = ctypes.string_at(text).decode("utf-8","replace").strip()
return comgr.AMD_COMGR_STATUS_SUCCESS
# nop callback
@comgr.amd_comgr_create_disassembly_info.argtypes[3]
def addr_cb(*args): return comgr.AMD_COMGR_STATUS_SUCCESS
def comgr_get_address_table(lib:bytes) -> dict[int, tuple[str, int]]:
check(comgr.amd_comgr_create_data(comgr.AMD_COMGR_DATA_KIND_EXECUTABLE, ctypes.byref(data_src:=comgr.amd_comgr_data_t())))
lib_buf = ctypes.create_string_buffer(lib, len(lib))
check(comgr.amd_comgr_set_data(data_src, len(lib), lib_buf))
check(comgr.amd_comgr_get_data_isa_name(data_src, isa_sz:=ctypes.c_size_t(128), isa:=(ctypes.c_char*isa_sz.value)()))
@comgr.amd_comgr_create_disassembly_info.argtypes[1]
def memory_cb(from_addr, to, size, _):
base, buf_len = ctypes.addressof(lib_buf), len(lib_buf)
start = int(from_addr) - base
if start < 0 or start >= buf_len: return 0
ctypes.memmove(to, base + start, n:=min(int(size), buf_len - start))
return n
info_src = comgr.amd_comgr_disassembly_info_t()
check(comgr.amd_comgr_create_disassembly_info(ctypes.cast(isa, ctypes.POINTER(ctypes.c_char)), memory_cb, instr_cb, addr_cb, info_src))
@comgr.amd_comgr_iterate_symbols.argtypes[1]
def sym_callback(sym, udata):
check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_TYPE, ctypes.byref(sym_type:=ctypes.c_int())))
if sym_type.value != comgr.AMD_COMGR_SYMBOL_TYPE_FUNC: return comgr.AMD_COMGR_STATUS_SUCCESS
check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_VALUE, ctypes.byref(vaddr:=ctypes.c_uint64())))
check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_SIZE, ctypes.byref(size:=ctypes.c_uint64())))
check(comgr.amd_comgr_map_elf_virtual_address_to_code_object_offset(data_src, vaddr.value, ctypes.byref(offset:=ctypes.c_uint64()),
ctypes.byref(ctypes.c_uint64()), ctypes.byref(nobits:=ctypes.c_bool())))
check(nobits.value)
base = ctypes.addressof(lib_buf)
pc = base + offset.value
end = pc + size.value
addr_table = ctypes.cast(udata, ctypes.POINTER(ctypes.py_object)).contents.value
instr_ref = ctypes.py_object(ctx:=InstrCtx())
instr_ptr = ctypes.cast(ctypes.pointer(instr_ref), ctypes.c_void_p)
while pc < end:
size_read = ctypes.c_uint64(0)
ctx.pc = pc
st = comgr.amd_comgr_disassemble_instruction(info_src, ctypes.c_uint64(pc), instr_ptr, ctypes.byref(size_read))
if st == comgr.AMD_COMGR_STATUS_SUCCESS and size_read.value:
rel = (pc - base) - offset.value
addr_table[vaddr.value + rel] = (ctx.inst, int(size_read.value))
pc += size_read.value
else: # don't inf loop if comgr fails
b = ctypes.c_ubyte.from_buffer(lib_buf, pc - base).value
addr_table[vaddr.value + (pc - base - offset.value)] = (f"DISASSEMBLER ISSUE 0x{b:02x}", 1)
pc += 1
return comgr.AMD_COMGR_STATUS_SUCCESS
addr_table:dict[int, tuple[str, int]] = {}
check(comgr.amd_comgr_iterate_symbols(data_src, sym_callback, ctypes.cast(ctypes.pointer(ctypes.py_object(addr_table)), ctypes.c_void_p)))
return addr_table
+3 -6
View File
@@ -156,9 +156,6 @@ class RGP:
sqtt_events = [x for x in profile if isinstance(x, ProfileSQTTEvent) and x.device == device_event.device]
if len(sqtt_events) == 0: raise RuntimeError(f"Device {device_event.device} doesn't contain SQTT data")
device_props = sqtt_events[0].props
gfx_ver = device_props['gfx_target_version'] // 10000
gfx_iplvl = getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}_{(device_props['gfx_target_version']//100)%100}",
getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}", None))
sqtt_itrace_enabled = any([event.itrace for event in sqtt_events])
sqtt_itrace_masked = not all_same([event.itrace for event in sqtt_events])
sqtt_itrace_se_mask = functools.reduce(lambda a,b: a|b, [int(event.itrace) << event.se for event in sqtt_events], 0) if sqtt_itrace_masked else 0
@@ -196,7 +193,7 @@ class RGP:
flags=0,
trace_shader_core_clock=0x93f05080,
trace_memory_clock=0x4a723a40,
device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550}[device_props['gfx_target_version']],
device_id={110000: 0x744c, 110003: 0x7480}[device_props['gfx_target_version']],
device_revision_id=0xc8,
vgprs_per_simd=1536,
sgprs_per_simd=128*16,
@@ -210,7 +207,7 @@ class RGP:
sgpr_alloc_granularity=128,
hardware_contexts=8,
gpu_type=sqtt.SQTT_GPU_TYPE_DISCRETE,
gfxip_level=gfx_iplvl,
gfxip_level=sqtt.SQTT_GFXIP_LEVEL_GFXIP_11_0,
gpu_index=0,
gds_size=0,
gds_per_shader_engine=0,
@@ -261,7 +258,7 @@ class RGP:
major_version=0, minor_version=2,
),
shader_engine_index=sqtt_event.se,
sqtt_version={11: sqtt.SQTT_VERSION_3_2, 12: sqtt.SQTT_VERSION_3_3}.get(gfx_ver),
sqtt_version=sqtt.SQTT_VERSION_3_2,
_0=sqtt.union_sqtt_file_chunk_sqtt_desc_0(
v1=sqtt.struct_sqtt_file_chunk_sqtt_desc_0_v1(
instrumentation_spec_version=1,
-96
View File
@@ -1,96 +0,0 @@
import ctypes, pathlib, argparse, pickle, re, functools, dataclasses
from extra.sqtt.rocprof import rocprof
from extra.sqtt.disasm import comgr_get_address_table
from tinygrad.helpers import temp, DEBUG
from tinygrad.device import ProfileEvent, ProfileProgramEvent
from tinygrad.runtime.ops_amd import ProfileSQTTEvent
@dataclasses.dataclass
class InstInfo:
typ:str=""
inst:str=""
hit:int=0
lat:int=0
stall:int=0
def __str__(self): return f"{self.inst:>20} hits:{self.typ:>6} hits:{self.hit:>6} latency:{self.lat:>6} stall:{self.stall:>6}"
def on_ev(self, ev):
self.hit, self.lat, self.stall = self.hit + 1, self.lat + ev.duration, self.stall + ev.stall
class _ROCParseCtx:
def __init__(self, sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
self.sqtt_evs, self.prog_evs = iter(sqtt_evs), prog_evs
self.wave_events, self.disasms, self.addr2prg = {}, {}, {}
for prog in prog_evs:
for addr, info in comgr_get_address_table(prog.lib).items():
self.disasms[prog.base + addr] = info
self.addr2prg[prog.base + addr] = prog
def next_sqtt(self): return next(self.sqtt_evs, None)
def find_program(self, addr): return self.addr2prg[addr]
def on_occupancy_ev(self, ev):
if DEBUG >= 4: print("OCC", ev.time, ev.cu, ev.simd, ev.wave_id, ev.start)
def on_wave_ev(self, ev):
if DEBUG >= 4: print("WAVE", ev.wave_id, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
asm = {}
for j in range(ev.instructions_size):
inst_ev = ev.instructions_array[j]
inst_typ = rocprof.rocprofiler_thread_trace_decoder_inst_category_t__enumvalues[inst_ev.category]
asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=self.disasms[inst_ev.pc.address][0]))
asm[inst_ev.pc.address].on_ev(inst_ev)
self.wave_events[(self.find_program(ev.instructions_array[0].pc.address).name, ev.wave_id, ev.cu, ev.simd)] = asm
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
args = parser.parse_args()
with args.profile.open("rb") as f: profile = pickle.load(f)
sqtt_events:list[ProfileSQTTEvent] = []
prog_events:list[ProfileProgramEvent] = []
for e in profile:
if isinstance(e, ProfileSQTTEvent): sqtt_events.append(e)
if isinstance(e, ProfileProgramEvent) and e.device.startswith("AMD"): prog_events.append(e)
ROCParseCtx = _ROCParseCtx(sqtt_events, prog_events)
@rocprof.rocprof_trace_decoder_se_data_callback_t
def copy_cb(buf, buf_size, data_ptr):
if (prof:=ROCParseCtx.next_sqtt()) is None: return 0
buf[0] = ctypes.cast((ctypes.c_ubyte * len(prof.blob)).from_buffer_copy(prof.blob), ctypes.POINTER(ctypes.c_ubyte))
buf_size[0] = len(prof.blob)
return len(prof.blob)
@rocprof.rocprof_trace_decoder_trace_callback_t
def trace_cb(record_type, events_ptr, n, data_ptr):
match record_type:
case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr): ROCParseCtx.on_occupancy_ev(ev)
case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr): ROCParseCtx.on_wave_ev(ev)
case _:
if DEBUG >= 2: print(rocprof.rocprofiler_thread_trace_decoder_record_type_t__enumvalues[record_type], events_ptr, n)
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
@rocprof.rocprof_trace_decoder_isa_callback_t
def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
instr, mem_size_ptr[0] = ROCParseCtx.disasms[pc.address]
# this is the number of bytes to next instruction, set to 0 for end_pgm
if instr == "s_endpgm": mem_size_ptr[0] = 0
if (max_sz:=size_ptr[0]) == 0: return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES
# truncate the instr if it doesn't fit
if (str_sz:=len(instr_bytes:=instr.encode()))+1 > max_sz: str_sz = max_sz
ctypes.memmove(instr_ptr, instr_bytes, str_sz)
size_ptr[0] = str_sz
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
print(ROCParseCtx.wave_events.keys())
-18
View File
@@ -1,18 +0,0 @@
#!/usr/bin/env python3
import os, shutil
from pathlib import Path
from tinygrad.helpers import fetch, OSX
DEST = Path("/usr/local/lib")
DEST.mkdir(exist_ok=True)
if __name__ == "__main__":
if OSX:
fp = fetch("https://github.com/ROCm/rocprof-trace-decoder/releases/download/0.1.4/rocprof-trace-decoder-macos-arm64-0.1.4-Darwin.sh")
lib = fp.parent/"rocprof-trace-decoder-macos-arm64-0.1.4-Darwin"/"lib"/"librocprof-trace-decoder.dylib"
os.chmod(fp, 0o755)
os.system(f"sudo {fp} --prefix={fp.parent} --include-subdir")
else:
lib = fetch("https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so", name="librocprof-trace-decoder.so")
shutil.copy2(lib, DEST)
print(f"Installed {lib.name} to", DEST)
-656
View File
@@ -1,656 +0,0 @@
# pylint: skip-file
# mypy: ignore-errors
# -*- coding: utf-8 -*-
#
# TARGET arch is: []
# WORD_SIZE is: 8
# POINTER_SIZE is: 8
# LONGDOUBLE_SIZE is: 16
#
import ctypes, ctypes.util
class AsDictMixin:
@classmethod
def as_dict(cls, self):
result = {}
if not isinstance(self, AsDictMixin):
# not a structure, assume it's already a python object
return self
if not hasattr(cls, "_fields_"):
return result
# sys.version_info >= (3, 5)
# for (field, *_) in cls._fields_: # noqa
for field_tuple in cls._fields_: # noqa
field = field_tuple[0]
if field.startswith('PADDING_'):
continue
value = getattr(self, field)
type_ = type(value)
if hasattr(value, "_length_") and hasattr(value, "_type_"):
# array
if not hasattr(type_, "as_dict"):
value = [v for v in value]
else:
type_ = type_._type_
value = [type_.as_dict(v) for v in value]
elif hasattr(value, "contents") and hasattr(value, "_type_"):
# pointer
try:
if not hasattr(type_, "as_dict"):
value = value.contents
else:
type_ = type_._type_
value = type_.as_dict(value.contents)
except ValueError:
# nullptr
value = None
elif isinstance(value, AsDictMixin):
# other structure
value = type_.as_dict(value)
result[field] = value
return result
class Structure(ctypes.Structure, AsDictMixin):
def __init__(self, *args, **kwds):
# We don't want to use positional arguments fill PADDING_* fields
args = dict(zip(self.__class__._field_names_(), args))
args.update(kwds)
super(Structure, self).__init__(**args)
@classmethod
def _field_names_(cls):
if hasattr(cls, '_fields_'):
return (f[0] for f in cls._fields_ if not f[0].startswith('PADDING'))
else:
return ()
@classmethod
def get_type(cls, field):
for f in cls._fields_:
if f[0] == field:
return f[1]
return None
@classmethod
def bind(cls, bound_fields):
fields = {}
for name, type_ in cls._fields_:
if hasattr(type_, "restype"):
if name in bound_fields:
if bound_fields[name] is None:
fields[name] = type_()
else:
# use a closure to capture the callback from the loop scope
fields[name] = (
type_((lambda callback: lambda *args: callback(*args))(
bound_fields[name]))
)
del bound_fields[name]
else:
# default callback implementation (does nothing)
try:
default_ = type_(0).restype().value
except TypeError:
default_ = None
fields[name] = type_((
lambda default_: lambda *args: default_)(default_))
else:
# not a callback function, use default initialization
if name in bound_fields:
fields[name] = bound_fields[name]
del bound_fields[name]
else:
fields[name] = type_()
if len(bound_fields) != 0:
raise ValueError(
"Cannot bind the following unknown callback(s) {}.{}".format(
cls.__name__, bound_fields.keys()
))
return cls(**fields)
class Union(ctypes.Union, AsDictMixin):
pass
c_int128 = ctypes.c_ubyte*16
c_uint128 = c_int128
void = None
if ctypes.sizeof(ctypes.c_longdouble) == 16:
c_long_double_t = ctypes.c_longdouble
else:
c_long_double_t = ctypes.c_ubyte*16
def string_cast(char_pointer, encoding='utf-8', errors='strict'):
value = ctypes.cast(char_pointer, ctypes.c_char_p).value
if value is not None and encoding is not None:
value = value.decode(encoding, errors=errors)
return value
def char_pointer_cast(string, encoding='utf-8'):
if encoding is not None:
try:
string = string.encode(encoding)
except AttributeError:
# In Python3, bytes has no encode attribute
pass
string = ctypes.c_char_p(string)
return ctypes.cast(string, ctypes.POINTER(ctypes.c_char))
class FunctionFactoryStub:
def __getattr__(self, _):
return ctypes.CFUNCTYPE(lambda y:y)
# libraries['FIXME_STUB'] explanation
# As you did not list (-l libraryname.so) a library that exports this function
# This is a non-working stub instead.
# You can either re-run clan2py with -l /path/to/library.so
# Or manually fix this by comment the ctypes.CDLL loading
_libraries = {}
_libraries['FIXME_STUB'] = ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder')) # ctypes.CDLL('FIXME_STUB')
# values for enumeration 'rocprofiler_thread_trace_decoder_info_t'
rocprofiler_thread_trace_decoder_info_t__enumvalues = {
0: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE',
1: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST',
2: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE',
3: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE',
4: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST',
}
ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE = 0
ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST = 1
ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE = 2
ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE = 3
ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST = 4
rocprofiler_thread_trace_decoder_info_t = ctypes.c_uint32 # enum
class struct_rocprofiler_thread_trace_decoder_pc_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_pc_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_pc_t._fields_ = [
('address', ctypes.c_uint64),
('code_object_id', ctypes.c_uint64),
]
rocprofiler_thread_trace_decoder_pc_t = struct_rocprofiler_thread_trace_decoder_pc_t
class struct_rocprofiler_thread_trace_decoder_perfevent_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_perfevent_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_perfevent_t._fields_ = [
('time', ctypes.c_int64),
('events0', ctypes.c_uint16),
('events1', ctypes.c_uint16),
('events2', ctypes.c_uint16),
('events3', ctypes.c_uint16),
('CU', ctypes.c_ubyte),
('bank', ctypes.c_ubyte),
('PADDING_0', ctypes.c_ubyte * 6),
]
rocprofiler_thread_trace_decoder_perfevent_t = struct_rocprofiler_thread_trace_decoder_perfevent_t
class struct_rocprofiler_thread_trace_decoder_occupancy_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_occupancy_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_occupancy_t._fields_ = [
('pc', rocprofiler_thread_trace_decoder_pc_t),
('time', ctypes.c_uint64),
('reserved', ctypes.c_ubyte),
('cu', ctypes.c_ubyte),
('simd', ctypes.c_ubyte),
('wave_id', ctypes.c_ubyte),
('start', ctypes.c_uint32, 1),
('_rsvd', ctypes.c_uint32, 31),
]
rocprofiler_thread_trace_decoder_occupancy_t = struct_rocprofiler_thread_trace_decoder_occupancy_t
# values for enumeration 'rocprofiler_thread_trace_decoder_wstate_type_t'
rocprofiler_thread_trace_decoder_wstate_type_t__enumvalues = {
0: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY',
1: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE',
2: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC',
3: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT',
4: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL',
5: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST',
}
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY = 0
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE = 1
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC = 2
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT = 3
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL = 4
ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST = 5
rocprofiler_thread_trace_decoder_wstate_type_t = ctypes.c_uint32 # enum
class struct_rocprofiler_thread_trace_decoder_wave_state_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_wave_state_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_wave_state_t._fields_ = [
('type', ctypes.c_int32),
('duration', ctypes.c_int32),
]
rocprofiler_thread_trace_decoder_wave_state_t = struct_rocprofiler_thread_trace_decoder_wave_state_t
# values for enumeration 'rocprofiler_thread_trace_decoder_inst_category_t'
rocprofiler_thread_trace_decoder_inst_category_t__enumvalues = {
0: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE',
1: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM',
2: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU',
3: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM',
4: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT',
5: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS',
6: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU',
7: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP',
8: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT',
9: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED',
10: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT',
11: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE',
12: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH',
13: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST',
}
ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE = 0
ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM = 1
ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU = 2
ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM = 3
ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT = 4
ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS = 5
ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU = 6
ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP = 7
ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT = 8
ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED = 9
ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT = 10
ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE = 11
ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH = 12
ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST = 13
rocprofiler_thread_trace_decoder_inst_category_t = ctypes.c_uint32 # enum
class struct_rocprofiler_thread_trace_decoder_inst_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_inst_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_inst_t._fields_ = [
('category', ctypes.c_uint32, 8),
('stall', ctypes.c_uint32, 24),
('duration', ctypes.c_int32),
('time', ctypes.c_int64),
('pc', rocprofiler_thread_trace_decoder_pc_t),
]
rocprofiler_thread_trace_decoder_inst_t = struct_rocprofiler_thread_trace_decoder_inst_t
class struct_rocprofiler_thread_trace_decoder_wave_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_wave_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_wave_t._fields_ = [
('cu', ctypes.c_ubyte),
('simd', ctypes.c_ubyte),
('wave_id', ctypes.c_ubyte),
('contexts', ctypes.c_ubyte),
('_rsvd1', ctypes.c_uint32),
('_rsvd2', ctypes.c_uint32),
('_rsvd3', ctypes.c_uint32),
('begin_time', ctypes.c_int64),
('end_time', ctypes.c_int64),
('timeline_size', ctypes.c_uint64),
('instructions_size', ctypes.c_uint64),
('timeline_array', ctypes.POINTER(struct_rocprofiler_thread_trace_decoder_wave_state_t)),
('instructions_array', ctypes.POINTER(struct_rocprofiler_thread_trace_decoder_inst_t)),
]
rocprofiler_thread_trace_decoder_wave_t = struct_rocprofiler_thread_trace_decoder_wave_t
class struct_rocprofiler_thread_trace_decoder_realtime_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_realtime_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_realtime_t._fields_ = [
('shader_clock', ctypes.c_int64),
('realtime_clock', ctypes.c_uint64),
('reserved', ctypes.c_uint64),
]
rocprofiler_thread_trace_decoder_realtime_t = struct_rocprofiler_thread_trace_decoder_realtime_t
# values for enumeration 'rocprofiler_thread_trace_decoder_shaderdata_flags_t'
rocprofiler_thread_trace_decoder_shaderdata_flags_t__enumvalues = {
0: 'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM',
1: 'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV',
}
ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM = 0
ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV = 1
rocprofiler_thread_trace_decoder_shaderdata_flags_t = ctypes.c_uint32 # enum
class struct_rocprofiler_thread_trace_decoder_shaderdata_t(Structure):
pass
struct_rocprofiler_thread_trace_decoder_shaderdata_t._pack_ = 1 # source:False
struct_rocprofiler_thread_trace_decoder_shaderdata_t._fields_ = [
('time', ctypes.c_int64),
('value', ctypes.c_uint64),
('cu', ctypes.c_ubyte),
('simd', ctypes.c_ubyte),
('wave_id', ctypes.c_ubyte),
('flags', ctypes.c_ubyte),
('reserved', ctypes.c_uint32),
]
rocprofiler_thread_trace_decoder_shaderdata_t = struct_rocprofiler_thread_trace_decoder_shaderdata_t
# values for enumeration 'rocprofiler_thread_trace_decoder_record_type_t'
rocprofiler_thread_trace_decoder_record_type_t__enumvalues = {
0: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP',
1: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY',
2: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT',
3: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE',
4: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO',
5: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG',
6: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA',
7: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME',
8: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY',
9: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST',
}
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP = 0
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY = 1
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT = 2
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE = 3
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO = 4
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG = 5
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA = 6
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME = 7
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY = 8
ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST = 9
rocprofiler_thread_trace_decoder_record_type_t = ctypes.c_uint32 # enum
# values for enumeration 'c__EA_rocprofiler_thread_trace_decoder_status_t'
c__EA_rocprofiler_thread_trace_decoder_status_t__enumvalues = {
0: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS',
1: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR',
2: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES',
3: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT',
4: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA',
5: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST',
}
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS = 0
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR = 1
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES = 2
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT = 3
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA = 4
ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST = 5
c__EA_rocprofiler_thread_trace_decoder_status_t = ctypes.c_uint32 # enum
rocprofiler_thread_trace_decoder_status_t = c__EA_rocprofiler_thread_trace_decoder_status_t
rocprofiler_thread_trace_decoder_status_t__enumvalues = c__EA_rocprofiler_thread_trace_decoder_status_t__enumvalues
rocprof_trace_decoder_trace_callback_t = ctypes.CFUNCTYPE(c__EA_rocprofiler_thread_trace_decoder_status_t, rocprofiler_thread_trace_decoder_record_type_t, ctypes.POINTER(None), ctypes.c_uint64, ctypes.POINTER(None))
rocprof_trace_decoder_isa_callback_t = ctypes.CFUNCTYPE(c__EA_rocprofiler_thread_trace_decoder_status_t, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint64), struct_rocprofiler_thread_trace_decoder_pc_t, ctypes.POINTER(None))
rocprof_trace_decoder_se_data_callback_t = ctypes.CFUNCTYPE(ctypes.c_uint64, ctypes.POINTER(ctypes.POINTER(ctypes.c_ubyte)), ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(None))
try:
rocprof_trace_decoder_parse_data = _libraries['FIXME_STUB'].rocprof_trace_decoder_parse_data
rocprof_trace_decoder_parse_data.restype = rocprofiler_thread_trace_decoder_status_t
rocprof_trace_decoder_parse_data.argtypes = [rocprof_trace_decoder_se_data_callback_t, rocprof_trace_decoder_trace_callback_t, rocprof_trace_decoder_isa_callback_t, ctypes.POINTER(None)]
except AttributeError:
pass
try:
rocprof_trace_decoder_get_info_string = _libraries['FIXME_STUB'].rocprof_trace_decoder_get_info_string
rocprof_trace_decoder_get_info_string.restype = ctypes.POINTER(ctypes.c_char)
rocprof_trace_decoder_get_info_string.argtypes = [rocprofiler_thread_trace_decoder_info_t]
except AttributeError:
pass
try:
rocprof_trace_decoder_get_status_string = _libraries['FIXME_STUB'].rocprof_trace_decoder_get_status_string
rocprof_trace_decoder_get_status_string.restype = ctypes.POINTER(ctypes.c_char)
rocprof_trace_decoder_get_status_string.argtypes = [rocprofiler_thread_trace_decoder_status_t]
except AttributeError:
pass
rocprofiler_thread_trace_decoder_debug_callback_t = ctypes.CFUNCTYPE(None, ctypes.c_int64, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), ctypes.POINTER(None))
uint64_t = ctypes.c_uint64
try:
rocprof_trace_decoder_dump_data = _libraries['FIXME_STUB'].rocprof_trace_decoder_dump_data
rocprof_trace_decoder_dump_data.restype = rocprofiler_thread_trace_decoder_status_t
rocprof_trace_decoder_dump_data.argtypes = [ctypes.POINTER(ctypes.c_char), uint64_t, rocprofiler_thread_trace_decoder_debug_callback_t, ctypes.POINTER(None)]
except AttributeError:
pass
class union_rocprof_trace_decoder_gfx9_header_t(Union):
pass
class struct_rocprof_trace_decoder_gfx9_header_t_0(Structure):
pass
struct_rocprof_trace_decoder_gfx9_header_t_0._pack_ = 1 # source:False
struct_rocprof_trace_decoder_gfx9_header_t_0._fields_ = [
('legacy_version', ctypes.c_uint64, 13),
('gfx9_version2', ctypes.c_uint64, 3),
('DSIMDM', ctypes.c_uint64, 4),
('DCU', ctypes.c_uint64, 5),
('reserved1', ctypes.c_uint64, 1),
('SEID', ctypes.c_uint64, 6),
('reserved2', ctypes.c_uint64, 32),
]
union_rocprof_trace_decoder_gfx9_header_t._pack_ = 1 # source:False
union_rocprof_trace_decoder_gfx9_header_t._anonymous_ = ('_0',)
union_rocprof_trace_decoder_gfx9_header_t._fields_ = [
('_0', struct_rocprof_trace_decoder_gfx9_header_t_0),
('raw', ctypes.c_uint64),
]
rocprof_trace_decoder_gfx9_header_t = union_rocprof_trace_decoder_gfx9_header_t
class union_rocprof_trace_decoder_instrument_enable_t(Union):
pass
class struct_rocprof_trace_decoder_instrument_enable_t_0(Structure):
pass
struct_rocprof_trace_decoder_instrument_enable_t_0._pack_ = 1 # source:False
struct_rocprof_trace_decoder_instrument_enable_t_0._fields_ = [
('char1', ctypes.c_uint32, 8),
('char2', ctypes.c_uint32, 8),
('char3', ctypes.c_uint32, 8),
('char4', ctypes.c_uint32, 8),
]
union_rocprof_trace_decoder_instrument_enable_t._pack_ = 1 # source:False
union_rocprof_trace_decoder_instrument_enable_t._anonymous_ = ('_0',)
union_rocprof_trace_decoder_instrument_enable_t._fields_ = [
('_0', struct_rocprof_trace_decoder_instrument_enable_t_0),
('u32All', ctypes.c_uint32),
]
rocprof_trace_decoder_instrument_enable_t = union_rocprof_trace_decoder_instrument_enable_t
class union_rocprof_trace_decoder_packet_header_t(Union):
pass
class struct_rocprof_trace_decoder_packet_header_t_0(Structure):
pass
struct_rocprof_trace_decoder_packet_header_t_0._pack_ = 1 # source:False
struct_rocprof_trace_decoder_packet_header_t_0._fields_ = [
('opcode', ctypes.c_uint32, 8),
('type', ctypes.c_uint32, 4),
('data20', ctypes.c_uint32, 20),
]
union_rocprof_trace_decoder_packet_header_t._pack_ = 1 # source:False
union_rocprof_trace_decoder_packet_header_t._anonymous_ = ('_0',)
union_rocprof_trace_decoder_packet_header_t._fields_ = [
('_0', struct_rocprof_trace_decoder_packet_header_t_0),
('u32All', ctypes.c_uint32),
]
rocprof_trace_decoder_packet_header_t = union_rocprof_trace_decoder_packet_header_t
# values for enumeration 'rocprof_trace_decoder_packet_opcode_t'
rocprof_trace_decoder_packet_opcode_t__enumvalues = {
4: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ',
5: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP',
6: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO',
}
ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ = 4
ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP = 5
ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO = 6
rocprof_trace_decoder_packet_opcode_t = ctypes.c_uint32 # enum
# values for enumeration 'rocprof_trace_decoder_agent_info_type_t'
rocprof_trace_decoder_agent_info_type_t__enumvalues = {
0: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ',
1: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL',
2: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST',
}
ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ = 0
ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL = 1
ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST = 2
rocprof_trace_decoder_agent_info_type_t = ctypes.c_uint32 # enum
class union_rocprof_trace_decoder_codeobj_marker_tail_t(Union):
pass
class struct_rocprof_trace_decoder_codeobj_marker_tail_t_0(Structure):
pass
struct_rocprof_trace_decoder_codeobj_marker_tail_t_0._pack_ = 1 # source:False
struct_rocprof_trace_decoder_codeobj_marker_tail_t_0._fields_ = [
('isUnload', ctypes.c_uint32, 1),
('bFromStart', ctypes.c_uint32, 1),
('legacy_id', ctypes.c_uint32, 30),
]
union_rocprof_trace_decoder_codeobj_marker_tail_t._pack_ = 1 # source:False
union_rocprof_trace_decoder_codeobj_marker_tail_t._anonymous_ = ('_0',)
union_rocprof_trace_decoder_codeobj_marker_tail_t._fields_ = [
('_0', struct_rocprof_trace_decoder_codeobj_marker_tail_t_0),
('raw', ctypes.c_uint32),
]
rocprof_trace_decoder_codeobj_marker_tail_t = union_rocprof_trace_decoder_codeobj_marker_tail_t
# values for enumeration 'rocprof_trace_decoder_codeobj_marker_type_t'
rocprof_trace_decoder_codeobj_marker_type_t__enumvalues = {
0: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL',
1: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO',
2: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO',
3: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI',
4: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI',
5: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO',
6: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI',
7: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST',
}
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL = 0
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO = 1
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO = 2
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI = 3
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI = 4
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO = 5
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI = 6
ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST = 7
rocprof_trace_decoder_codeobj_marker_type_t = ctypes.c_uint32 # enum
__all__ = \
['ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST',
'ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST',
'ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE',
'ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE',
'ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU',
'ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA',
'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE',
'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM',
'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV',
'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR',
'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT',
'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA',
'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES',
'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST',
'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS',
'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY',
'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC',
'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE',
'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST',
'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL',
'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT',
'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL',
'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST',
'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO',
'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL',
'ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO',
'ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ',
'ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP',
'c__EA_rocprofiler_thread_trace_decoder_status_t',
'rocprof_trace_decoder_agent_info_type_t',
'rocprof_trace_decoder_codeobj_marker_tail_t',
'rocprof_trace_decoder_codeobj_marker_type_t',
'rocprof_trace_decoder_dump_data',
'rocprof_trace_decoder_get_info_string',
'rocprof_trace_decoder_get_status_string',
'rocprof_trace_decoder_gfx9_header_t',
'rocprof_trace_decoder_instrument_enable_t',
'rocprof_trace_decoder_isa_callback_t',
'rocprof_trace_decoder_packet_header_t',
'rocprof_trace_decoder_packet_opcode_t',
'rocprof_trace_decoder_parse_data',
'rocprof_trace_decoder_se_data_callback_t',
'rocprof_trace_decoder_trace_callback_t',
'rocprofiler_thread_trace_decoder_debug_callback_t',
'rocprofiler_thread_trace_decoder_info_t',
'rocprofiler_thread_trace_decoder_inst_category_t',
'rocprofiler_thread_trace_decoder_inst_t',
'rocprofiler_thread_trace_decoder_occupancy_t',
'rocprofiler_thread_trace_decoder_pc_t',
'rocprofiler_thread_trace_decoder_perfevent_t',
'rocprofiler_thread_trace_decoder_realtime_t',
'rocprofiler_thread_trace_decoder_record_type_t',
'rocprofiler_thread_trace_decoder_shaderdata_flags_t',
'rocprofiler_thread_trace_decoder_shaderdata_t',
'rocprofiler_thread_trace_decoder_status_t',
'rocprofiler_thread_trace_decoder_status_t__enumvalues',
'rocprofiler_thread_trace_decoder_wave_state_t',
'rocprofiler_thread_trace_decoder_wave_t',
'rocprofiler_thread_trace_decoder_wstate_type_t',
'struct_rocprof_trace_decoder_codeobj_marker_tail_t_0',
'struct_rocprof_trace_decoder_gfx9_header_t_0',
'struct_rocprof_trace_decoder_instrument_enable_t_0',
'struct_rocprof_trace_decoder_packet_header_t_0',
'struct_rocprofiler_thread_trace_decoder_inst_t',
'struct_rocprofiler_thread_trace_decoder_occupancy_t',
'struct_rocprofiler_thread_trace_decoder_pc_t',
'struct_rocprofiler_thread_trace_decoder_perfevent_t',
'struct_rocprofiler_thread_trace_decoder_realtime_t',
'struct_rocprofiler_thread_trace_decoder_shaderdata_t',
'struct_rocprofiler_thread_trace_decoder_wave_state_t',
'struct_rocprofiler_thread_trace_decoder_wave_t', 'uint64_t',
'union_rocprof_trace_decoder_codeobj_marker_tail_t',
'union_rocprof_trace_decoder_gfx9_header_t',
'union_rocprof_trace_decoder_instrument_enable_t',
'union_rocprof_trace_decoder_packet_header_t']
-5
View File
@@ -43,7 +43,6 @@ enum sqtt_version
SQTT_VERSION_2_3 = 0x6, /* GFX9 */
SQTT_VERSION_2_4 = 0x7, /* GFX10+ */
SQTT_VERSION_3_2 = 0xb, /* GFX11+ */
SQTT_VERSION_3_3 = 0xc, /* GFX12+ */
};
enum sqtt_file_chunk_type
@@ -145,8 +144,6 @@ enum sqtt_gfxip_level
SQTT_GFXIP_LEVEL_GFXIP_10_1 = 0x7,
SQTT_GFXIP_LEVEL_GFXIP_10_3 = 0x9,
SQTT_GFXIP_LEVEL_GFXIP_11_0 = 0xc,
SQTT_GFXIP_LEVEL_GFXIP_11_5 = 0xd,
SQTT_GFXIP_LEVEL_GFXIP_12 = 0x10,
};
enum sqtt_memory_type
@@ -430,8 +427,6 @@ enum elf_gfxip_level
EF_AMDGPU_MACH_AMDGCN_GFX1010 = 0x033,
EF_AMDGPU_MACH_AMDGCN_GFX1030 = 0x036,
EF_AMDGPU_MACH_AMDGCN_GFX1100 = 0x041,
EF_AMDGPU_MACH_AMDGCN_GFX1150 = 0x043,
EF_AMDGPU_MACH_AMDGCN_GFX1200 = 0x04e,
};
struct sqtt_file_chunk_spm_db {
@@ -1,400 +0,0 @@
/**
* @file
* @brief Basic operations on generic types.
*/
#pragma once
#include <cuda_bf16.h>
#include <limits>
#include "base_types.cuh"
namespace kittens {
/**
* @namespace base_ops
*
* @brief A namespace for operations on basic data types.
*/
namespace base_ops {
/* ---------- CONST OPS ---------- */
/**
* @brief Represents the zero constant operation.
*
* This operation returns the zero value of the specified type.
*
* @tparam T The data type for which to return the zero value.
* @return The zero value of type T.
*/
struct zero {
template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::zero(); }
};
/**
* @brief Represents the one constant operation.
*
* This operation returns the one value of the specified type.
*
* @tparam T The data type for which to return the one value.
* @return The one value of type T.
*/
struct one {
template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::one(); }
};
/**
* @brief Represents the positive infinity constant operation.
*
* This operation returns the positive infinity value of the specified type.
*
* @tparam T The data type for which to return the positive infinity value.
* @return The positive infinity value of type T.
*/
struct pos_infty {
template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::pos_infty(); }
};
/**
* @brief Represents the negative infinity constant operation.
*
* This operation returns the negative infinity value of the specified type.
*
* @tparam T The data type for which to return the negative infinity value.
* @return The negative infinity value of type T.
*/
struct neg_infty {
template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::neg_infty(); }
};
/* ---------- UNARY OPS ---------- */
/**
* @brief Exponential function operation.
*
* This operation calculates the exponential of the input value.
*
* @tparam T The data type of the input and output values.
* @param x[in] The input value.
* @return The exponential of the input value.
*/
struct exp {
template<typename T> static __device__ inline T op(const T &x) { return exp(x); }
};
template<> __device__ inline float exp::op<float> (const float &x ) { return __expf(x); }
template<> __device__ inline float2 exp::op<float2>(const float2 &x) { return float2{__expf(x.x), __expf(x.y)}; }
template<> __device__ inline bf16 exp::op<bf16> (const bf16 &x ) { return hexp(x); }
template<> __device__ inline bf16_2 exp::op<bf16_2>(const bf16_2 &x) { return h2exp(x); }
template<> __device__ inline half exp::op<half> (const half &x ) { return hexp(x); }
template<> __device__ inline half_2 exp::op<half_2>(const half_2 &x) { return h2exp(x); }
/**
* @brief Exponential function operation, in base 2
*
* This operation calculates the exponential of the input value, in base 2.
*
* @tparam T The data type of the input and output values.
* @param x[in] The input value.
* @return The exponential of the input value.
*/
struct exp2 {
template<typename T> static __device__ inline T op(const T &x) { return exp2f(x); }
};
template<> __device__ inline float exp2::op<float> (const float &x ) { return exp2f(x); }
template<> __device__ inline float2 exp2::op<float2>(const float2 &x) { return float2{exp2f(x.x), exp2f(x.y)}; }
template<> __device__ inline bf16 exp2::op<bf16> (const bf16 &x ) { return hexp2(x); }
template<> __device__ inline bf16_2 exp2::op<bf16_2>(const bf16_2 &x) { return h2exp2(x); }
template<> __device__ inline half exp2::op<half> (const half &x ) { return hexp2(x); }
template<> __device__ inline half_2 exp2::op<half_2>(const half_2 &x) { return h2exp2(x); }
/**
* @brief Natural log function operation.
*
* This operation calculates the natural logarithm of the input value.
*
* @tparam T The data type of the input and output values.
* @param x[in] The input value.
* @return The natural logarithm of the input value.
*/
struct log {
template<typename T> static __device__ inline T op(const T &x) { return log(x); }
};
template<> __device__ inline float log::op<float> (const float &x ) { return __logf(x); }
template<> __device__ inline float2 log::op<float2>(const float2 &x) { return float2{__logf(x.x), __logf(x.y)}; }
template<> __device__ inline bf16 log::op<bf16> (const bf16 &x ) { return hlog(x); }
template<> __device__ inline bf16_2 log::op<bf16_2>(const bf16_2 &x) { return h2log(x); }
template<> __device__ inline half log::op<half> (const half &x ) { return hlog(x); }
template<> __device__ inline half_2 log::op<half_2>(const half_2 &x) { return h2log(x); }
/**
* @brief Logarithm base 2 operation.
*
* This operation calculates the logarithm base 2 of the input value.
*
* @tparam T The data type of the input and output values.
* @param x[in] The input value.
* @return The logarithm base 2 of the input value.
*/
struct log2 {
template<typename T> static __device__ inline T op(const T &x) { return log2(x); }
};
template<> __device__ inline float log2::op<float> (const float &x ) { return __log2f(x); }
template<> __device__ inline float2 log2::op<float2>(const float2 &x) { return float2{__log2f(x.x), __log2f(x.y)}; }
template<> __device__ inline bf16 log2::op<bf16> (const bf16 &x ) { return hlog2(x); }
template<> __device__ inline bf16_2 log2::op<bf16_2>(const bf16_2 &x) { return h2log2(x); }
template<> __device__ inline half log2::op<half> (const half &x ) { return hlog2(x); }
template<> __device__ inline half_2 log2::op<half_2>(const half_2 &x) { return h2log2(x); }
/**
* @brief Absolute value operation.
*
* This operation calculates the absolute value of the input.
*
* @tparam T The data type of the input and output values.
* @param x[in] The input value.
* @return The absolute value of the input.
*/
struct abs {
template<typename T> static __device__ inline T op(const T &x) { return abs(x); }
};
template<> __device__ inline float abs::op<float> (const float &x ) { return fabsf(x); }
template<> __device__ inline float2 abs::op<float2>(const float2 &x) { return float2{fabsf(x.x), fabsf(x.y)}; }
template<> __device__ inline bf16 abs::op<bf16> (const bf16 &x ) { return __habs(x); }
template<> __device__ inline bf16_2 abs::op<bf16_2>(const bf16_2 &x) { return __habs2(x); }
template<> __device__ inline half abs::op<half> (const half &x ) { return __habs(x); }
template<> __device__ inline half_2 abs::op<half_2>(const half_2 &x) { return __habs2(x); }
/**
* @brief Rectified Linear Unit (ReLU) operation.
*
* This operation applies the ReLU function to the input, which is the
* maximum of zero and the input value.
*
* @tparam T The data type of the input and output values.
* @param x[in] The input value.
* @return The result of ReLU function applied to the input.
*/
struct relu {
template<typename T> static __device__ inline T op(const T &x) { return max(x, base_types::constants<T>::zero()); }
};
template<> __device__ inline float relu::op<float> (const float &x ) { return max(x, 0.f); }
template<> __device__ inline float2 relu::op<float2>(const float2 &x) { return float2{max(x.x, 0.f), max(x.y, 0.f)}; }
template<> __device__ inline bf16 relu::op<bf16> (const bf16 &x ) { return __hmax(x, base_types::constants<bf16>::zero()); }
template<> __device__ inline bf16_2 relu::op<bf16_2>(const bf16_2 &x) { return __hmax2(x, base_types::constants<bf16_2>::zero()); }
template<> __device__ inline half relu::op<half> (const half &x ) { return __hmax(x, base_types::constants<half>::zero()); }
template<> __device__ inline half_2 relu::op<half_2>(const half_2 &x) { return __hmax2(x, base_types::constants<half_2>::zero()); }
/**
* @brief Copy operation.
*
* This operation returns the input value unchanged.
*
* @tparam T The data type of the input and output values.
* @param a[in] The input value.
* @return The same value as the input.
*/
struct copy { // for non-compile-time setters.
template<typename T> static __device__ inline T op(const T &a) { return a; }
};
/* ---------- BINARY OPS ---------- */
/**
* @brief Copy2 operation.
*
* This operation returns the second input value unchanged.
*
* @tparam T The data type of the input and output values.
* @param a[in] The first input value (ignored).
* @param b[in] The second input value.
* @return The same value as the second input.
*/
struct copy2 { // this turns out to be a slightly hacky op that makes some code cleaner :/
template<typename T> static __device__ inline T op(const T &a, const T &b) { return b; }
};
/**
* @brief Sum operation.
*
* This operation calculates the sum of two input values.
*
* @tparam T The data type of the input and output values.
* @param a[in] The first input value.
* @param b[in] The second input value.
* @return The sum of the input values.
*/
struct sum {
template<typename T> static __device__ inline T op(const T &a, const T &b) { return a+b; }
};
template<> __device__ inline float2 sum::op<float2>(const float2 &a, const float2 &b) {
#ifdef KITTENS_BLACKWELL
float2 c;
asm volatile("add.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
return c;
#else
return float2{a.x+b.x, a.y+b.y};
#endif
}
template<> __device__ inline bf16 sum::op<bf16> (const bf16 &a, const bf16 &b) { return __hadd(a, b); }
template<> __device__ inline bf16_2 sum::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hadd2(a, b); }
template<> __device__ inline half sum::op<half> (const half &a, const half &b) { return __hadd(a, b); }
template<> __device__ inline half_2 sum::op<half_2>(const half_2 &a, const half_2 &b) { return __hadd2(a, b); }
/**
* @brief Subtraction operation.
*
* This operation calculates the difference between two input values.
*
* @tparam T The data type of the input and output values.
* @param a[in] The first input value.
* @param b[in] The second input value.
* @return The difference between the input values.
*/
struct sub {
template<typename T> static __device__ inline T op(const T &a, const T &b) { return a-b; }
};
template<> __device__ inline float2 sub::op<float2>(const float2 &a, const float2 &b) {
#ifdef KITTENS_BLACKWELL
float2 c;
asm volatile("sub.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
return c;
#else
return float2{a.x-b.x, a.y-b.y};
#endif
}
template<> __device__ inline bf16 sub::op<bf16> (const bf16 &a, const bf16 &b) { return __hsub(a, b); }
template<> __device__ inline bf16_2 sub::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hsub2(a, b); }
template<> __device__ inline half sub::op<half> (const half &a, const half &b) { return __hsub(a, b); }
template<> __device__ inline half_2 sub::op<half_2>(const half_2 &a, const half_2 &b) { return __hsub2(a, b); }
/**
* @brief Multiplication operation.
*
* This operation calculates the product of two input values.
*
* @tparam T The data type of the input and output values.
* @param a[in] The first input value.
* @param b[in] The second input value.
* @return The product of the input values.
*/
struct mul {
template<typename T> static __device__ inline T op(const T &a, const T &b) { return a*b; }
};
template<> __device__ inline float2 mul::op<float2>(const float2 &a, const float2 &b) {
#ifdef KITTENS_BLACKWELL
float2 c;
asm volatile("mul.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
return c;
#else
return float2{a.x*b.x, a.y*b.y};
#endif
}
template<> __device__ inline bf16 mul::op<bf16> (const bf16 &a, const bf16 &b) { return __hmul(a, b); }
template<> __device__ inline bf16_2 mul::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmul2(a, b); }
template<> __device__ inline half mul::op<half> (const half &a, const half &b) { return __hmul(a, b); }
template<> __device__ inline half_2 mul::op<half_2>(const half_2 &a, const half_2 &b) { return __hmul2(a, b); }
/**
* @brief Division operation.
*
* This operation calculates the quotient of two input values.
*
* @tparam T The data type of the input and output values.
* @param a[in] The first input value.
* @param b[in] The second input value.
* @return The quotient of the input values.
*/
struct div {
template<typename T> static __device__ inline T op(const T &a, const T &b) { return a/b; }
};
template<> __device__ inline float2 div::op<float2>(const float2 &a, const float2 &b) { return float2{a.x/b.x, a.y/b.y}; }
template<> __device__ inline bf16 div::op<bf16> (const bf16 &a, const bf16 &b) { return __hdiv(a, b); }
template<> __device__ inline bf16_2 div::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __h2div(a, b); } // this op is a special snowflake
template<> __device__ inline half div::op<half> (const half &a, const half &b) { return __hdiv(a, b); }
template<> __device__ inline half_2 div::op<half_2>(const half_2 &a, const half_2 &b) { return __h2div(a, b); }
/**
* @brief Maximum operation.
*
* This operation calculates the maximum of two input values.
*
* @tparam T The data type of the input and output values.
* @param a[in] The first input value.
* @param b[in] The second input value.
* @return The maximum of the input values.
*/
struct max {
template<typename T> static __device__ inline T op(const T &a, const T &b) { return ::max(a, b); }
};
template<> __device__ inline float2 max::op<float2>(const float2 &a, const float2 &b) { return float2{::max(a.x, b.x), ::max(a.y, b.y)}; }
template<> __device__ inline bf16 max::op<bf16> (const bf16 &a, const bf16 &b) { return __hmax(a, b); }
template<> __device__ inline bf16_2 max::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmax2(a, b); }
template<> __device__ inline half max::op<half> (const half &a, const half &b) { return __hmax(a, b); }
template<> __device__ inline half_2 max::op<half_2>(const half_2 &a, const half_2 &b) { return __hmax2(a, b); }
/**
* @brief Minimum operation.
*
* This operation calculates the minimum of two input values.
*
* @tparam T The data type of the input and output values.
* @param a[in] The first input value.
* @param b[in] The second input value.
* @return The minimum of the input values.
*/
struct min {
template<typename T> static __device__ inline T op(const T &a, const T &b) { return ::min(a, b); }
};
template<> __device__ inline float2 min::op<float2>(const float2 &a, const float2 &b) { return float2{::min(a.x, b.x), ::min(a.y, b.y)}; }
template<> __device__ inline bf16 min::op<bf16> (const bf16 &a, const bf16 &b) { return __hmin(a, b); }
template<> __device__ inline bf16_2 min::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmin2(a, b); }
template<> __device__ inline half min::op<half> (const half &a, const half &b) { return __hmin(a, b); }
template<> __device__ inline half_2 min::op<half_2>(const half_2 &a, const half_2 &b) { return __hmin2(a, b); }
/* ---------- TERNARY OPS ---------- */
/**
* @brief Fused multiply-add operation A * B + C.
*
* This operation performs a fused multiply-add, computing (A * B) + C with only one rounding.
*
* @tparam T The data type of the input and output values.
* @param a[in] The first input value.
* @param b[in] The second input value.
* @param c[in] The third input value to be added.
* @return The result of the fused multiply-add operation.
*/
struct fma_AxBtC {
template<typename T> static __device__ inline T op(const T &a, const T &b, const T &c) {
return sum::op<T>(mul::op<T>(a, b), c);
}
};
template<> __device__ inline float2 fma_AxBtC::op<float2>(const float2 &a, const float2 &b, const float2 &c) {
#ifdef KITTENS_BLACKWELL
float2 d;
asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b), "l"(*(uint64_t*)&c));
return d;
#else
return float2{a.x*b.x+c.x, a.y*b.y+c.y};
#endif
}
/**
* @brief Fused multiply-add operation A * C + B.
*
* This operation performs a fused multiply-add, computing (A * C) + B with only one rounding.
* This is particularly useful for attention mechanisms in neural networks.
*
* @tparam T The data type of the input and output values.
* @param a[in] The first input value.
* @param b[in] The third input value to be added.
* @param c[in] The second input value.
* @return The result of the fused multiply-add operation.
*/
struct fma_AxCtB { // this is the one needed for attention
template<typename T> static __device__ inline T op(const T &a, const T &b, const T &c) {
return sum::op<T>(mul::op<T>(a, c), b);
}
};
template<> __device__ inline float2 fma_AxCtB::op<float2>(const float2 &a, const float2 &b, const float2 &c) {
#ifdef KITTENS_BLACKWELL
float2 d;
asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&c), "l"(*(uint64_t*)&b));
return d;
#else
return float2{a.x*c.x+b.x, a.y*c.y+b.y};
#endif
}
} // namespace base_ops
} // namespace kittens
@@ -1,519 +0,0 @@
/**
* @file
* @brief Declarations, manipulations, and wrappers for basic types.
*
* This file is a bunch of utilities for going back and forth between different types.
*
* Many of them are for the compiler, so as to clean up the code. It unfortunately
* seems necessary when we have types we really care about that are less than word width.
*/
#pragma once
#ifdef KITTENS_HOPPER
#include <cuda_fp8.h>
#endif
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <string>
#include <bit>
namespace kittens {
/**
* @brief Bfloat16 floating-point type.
*/
using bf16 = __nv_bfloat16;
/**
* @brief Half-precision floating-point type.
*/
using half = __half;
/**
* @brief Packed word of two bfloat16 floating-point values.
*/
using bf16_2 = __nv_bfloat162;
/**
* @brief Packed word of two half-precision floating-point values.
*/
using half_2 = __half2;
#ifdef KITTENS_HOPPER
/**
* @brief float8 floating-point type.
*/
using fp8e4m3 = __nv_fp8_e4m3;
using fp8e5m2 = __nv_fp8_e5m2;
#ifdef KITTENS_BLACKWELL
using fp8e8m0 = __nv_fp8_e8m0;
#endif
/**
* @brief 2-packed float8 floating-point type.
*/
using fp8e4m3_2 = __nv_fp8x2_e4m3;
using fp8e5m2_2 = __nv_fp8x2_e5m2;
#ifdef KITTENS_BLACKWELL
using fp8e8m0_2 = __nv_fp8x2_e8m0;
#endif
/**
* @brief 4-packed float8 floating-point type.
*/
using fp8e4m3_4 = __nv_fp8x4_e4m3;
using fp8e5m2_4 = __nv_fp8x4_e5m2;
#ifdef KITTENS_BLACKWELL
using fp8e8m0_4 = __nv_fp8x4_e8m0;
#endif
#endif
namespace ducks {
/**
* @namespace base_types
*
* @brief A namespace for concepts for basic data types.
*/
namespace base_types {
#ifdef KITTENS_HOPPER
#ifdef KITTENS_BLACKWELL
template<typename T>
concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2> || std::is_same_v<T, fp8e4m3_4> || std::is_same_v<T, fp8e5m2_4> || std::is_same_v<T, fp8e8m0_4>; // could add half_2 later if implemented.
template<typename T>
concept T1 = std::is_same_v<T, float> || std::is_same_v<T, bf16 > || std::is_same_v<T, half> || std::is_same_v<T, fp8e4m3> || std::is_same_v<T, fp8e5m2> || std::is_same_v<T, fp8e8m0>; // could add half_2 later if implemented.
#else
template<typename T>
concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2> || std::is_same_v<T, fp8e4m3_4> || std::is_same_v<T, fp8e5m2_4>;
template<typename T>
concept T1 = std::is_same_v<T, float> || std::is_same_v<T, bf16 > || std::is_same_v<T, half> || std::is_same_v<T, fp8e4m3> || std::is_same_v<T, fp8e5m2>;
#endif
#else
template<typename T>
concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2>;
template<typename T>
concept T1 = std::is_same_v<T, float> || std::is_same_v<T, bf16 > || std::is_same_v<T, half>;
#endif
} // namespace base_types
} // namespace ducks
/**
* @namespace base_types
*
* @brief A namespace for ThunderKittens basic data types.
*/
namespace base_types {
/**
* @brief Provides compile-time constants for different types.
*
* @tparam T The type for which to provide constants.
*/
template<typename T> struct constants {
/**
* @brief Zero
* @return Constexpr zero with type T
*/
static __device__ inline constexpr T zero() { return T{0}; }
/**
* @brief One
* @return Constexpr one with type T
*/
static __device__ inline constexpr T one() { return T{1}; }
/**
* @brief Positive infinity. Particularly useful for initializing before a min op.
* @return Constexpr positive infinity with type T
*/
static __device__ inline constexpr T pos_infty() { return T{INFINITY}; } // I'll find a better way at some point but this appears to work.
/**
* @brief Negative infinity. Particularly useful for initializing before a max op.
* @return Constexpr negative infinity with type T
*/
static __device__ inline constexpr T neg_infty() { return T{-INFINITY}; }
};
template<> struct constants<float2> {
static __device__ inline constexpr float2 zero() { return float2{0.f, 0.f}; }
static __device__ inline constexpr float2 one() { return float2{1.f, 1.f}; }
static __device__ inline constexpr float2 pos_infty() { return float2{constants<float>::pos_infty(), constants<float>::pos_infty()}; }
static __device__ inline constexpr float2 neg_infty() { return float2{constants<float>::neg_infty(), constants<float>::neg_infty()}; }
};
template<> struct constants<bf16> {
static __device__ inline constexpr bf16 zero() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x0000)); } // unfortunately __float2bf16_rn is not constexpr
static __device__ inline constexpr bf16 one() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x3F80)); }
static __device__ inline constexpr bf16 pos_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x7F80)); }
static __device__ inline constexpr bf16 neg_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0xFF80)); }
};
template<> struct constants<bf16_2> {
static __device__ inline constexpr bf16_2 zero() { return bf16_2{constants<bf16>::zero(), constants<bf16>::zero()}; }
static __device__ inline constexpr bf16_2 one() { return bf16_2{constants<bf16>::one(), constants<bf16>::one()}; }
static __device__ inline constexpr bf16_2 pos_infty() { return bf16_2{constants<bf16>::pos_infty(), constants<bf16>::pos_infty()}; }
static __device__ inline constexpr bf16_2 neg_infty() { return bf16_2{constants<bf16>::neg_infty(), constants<bf16>::neg_infty()}; }
};
template<> struct constants<half> {
static __device__ inline constexpr half zero() { return std::bit_cast<__half>(uint16_t(0x0000)); }
static __device__ inline constexpr half one() { return std::bit_cast<__half>(uint16_t(0x3C00)); }
static __device__ inline constexpr half pos_infty() { return std::bit_cast<__half>(uint16_t(0x7C00)); }
static __device__ inline constexpr half neg_infty() { return std::bit_cast<__half>(uint16_t(0xFC00)); }
};
template<> struct constants<half_2> {
static __device__ inline constexpr half_2 zero() { return half_2{constants<half>::zero(), constants<half>::zero()}; }
static __device__ inline constexpr half_2 one() { return half_2{constants<half>::one(), constants<half>::one()}; }
static __device__ inline constexpr half_2 pos_infty() { return half_2{constants<half>::pos_infty(), constants<half>::pos_infty()}; }
static __device__ inline constexpr half_2 neg_infty() { return half_2{constants<half>::neg_infty(), constants<half>::neg_infty()}; }
};
#ifdef KITTENS_HOPPER
template<> struct constants<fp8e4m3> {
static __device__ inline constexpr fp8e4m3 zero() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x00)); }
static __device__ inline constexpr fp8e4m3 one() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x38)); }
};
template<> struct constants<fp8e4m3_2> {
static __device__ inline constexpr fp8e4m3_2 zero() { return std::bit_cast<fp8e4m3_2>(uint16_t(0x0000)); }
static __device__ inline constexpr fp8e4m3_2 one() { return std::bit_cast<fp8e4m3_2>(uint16_t(0x3838)); }
};
template<> struct constants<fp8e4m3_4> {
static __device__ inline constexpr fp8e4m3_4 zero() { return std::bit_cast<fp8e4m3_4>(uint32_t(0x00000000)); }
static __device__ inline constexpr fp8e4m3_4 one() { return std::bit_cast<fp8e4m3_4>(uint32_t(0x38383838)); }
};
template<> struct constants<fp8e5m2> {
static __device__ inline constexpr fp8e5m2 zero() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x00)); }
static __device__ inline constexpr fp8e5m2 one() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x3C)); }
};
template<> struct constants<fp8e5m2_2> {
static __device__ inline constexpr fp8e5m2_2 zero() { return std::bit_cast<fp8e5m2_2>(uint16_t(0x0000)); }
static __device__ inline constexpr fp8e5m2_2 one() { return std::bit_cast<fp8e5m2_2>(uint16_t(0x3C3C)); }
};
template<> struct constants<fp8e5m2_4> {
static __device__ inline constexpr fp8e5m2_4 zero() { return std::bit_cast<fp8e5m2_4>(uint32_t(0x00000000)); }
static __device__ inline constexpr fp8e5m2_4 one() { return std::bit_cast<fp8e5m2_4>(uint32_t(0x3C3C3C3C)); }
};
#endif
template<> struct constants<int> {
static __device__ inline constexpr int zero() { return 0; }
static __device__ inline constexpr int one() { return 1; }
};
template<> struct constants<int2> {
static __device__ inline constexpr int2 zero() { return int2{0, 0}; }
static __device__ inline constexpr int2 one() { return int2{1, 1}; }
};
/**
* @brief Provides information about packing of elements for a given type.
*
* @tparam T The type for which to provide packing information.
*/
template<typename T> struct packing {
/**
* @brief The number of elements packed together.
*
* @return constexpr int representing number of elements within the type.
*/
static __device__ inline constexpr int num() { return 1; }
/**
* @brief Packs a single T element twice (replicated) into its packed type.
*
* @param i[in] The element to pack.
* @return The packed type.
*/
static __device__ inline constexpr T pack(const bf16 &i);
};
template<> struct packing<bf16> {
static __device__ inline constexpr int num() { return 1; }
using unpacked_type = bf16;
using packed_type = bf16_2;
static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; }
};
template<> struct packing<bf16_2> {
static __device__ inline constexpr int num() { return 2; }
using unpacked_type = bf16;
using packed_type = bf16_2;
static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; } // this replication makes code cleaner later.
};
template<> struct packing<half> {
static __device__ inline constexpr int num() { return 1; }
using unpacked_type = half;
using packed_type = half_2;
static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; }
};
template<> struct packing<half_2> {
static __device__ inline constexpr int num() { return 2; }
using unpacked_type = half;
using packed_type = half_2;
static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; } // this replication makes code cleaner later.
};
template<> struct packing<float> {
static __device__ inline constexpr int num() { return 1; }
using unpacked_type = float;
using packed_type = float2;
static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; }
};
template<> struct packing<float2> {
static __device__ inline constexpr int num() { return 2; }
using unpacked_type = float;
using packed_type = float2;
static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; } // this replication makes code cleaner later.
};
template<> struct packing<char> {
static __device__ inline constexpr int num() { return 1; }
using unpacked_type = char;
using packed_type = char2;
static __device__ inline constexpr char2 pack(const char &i) { return char2{i, i}; } // this replication makes code cleaner later.
};
template<> struct packing<char2> {
static __device__ inline constexpr int num() { return 2; }
using unpacked_type = char;
using packed_type = char2;
static __device__ inline constexpr char2 pack(const char &i) { return char2{i, i}; } // this replication makes code cleaner later.
};
template<> struct packing<int> {
static __device__ inline constexpr int num() { return 1; }
using unpacked_type = int;
using packed_type = int2;
static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } // this replication makes code cleaner later.
};
template<> struct packing<int2> {
static __device__ inline constexpr int num() { return 2; }
using unpacked_type = int;
using packed_type = int2;
static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } // this replication makes code cleaner later.
};
template<> struct packing<uint> {
static __device__ inline constexpr int num() { return 1; }
using unpacked_type = uint;
using packed_type = uint2;
static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } // this replication makes code cleaner later.
};
template<> struct packing<uint2> {
static __device__ inline constexpr int num() { return 2; }
using unpacked_type = uint;
using packed_type = uint2;
static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } // this replication makes code cleaner later.
};
struct uint64_2 { uint64_t x, y; };
template<> struct packing<uint64_t> {
static __device__ inline constexpr int num() { return 1; }
using unpacked_type = uint64_t;
using packed_type = uint64_2;
static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } // this replication makes code cleaner later.
};
template<> struct packing<uint64_2> {
static __device__ inline constexpr int num() { return 2; }
using unpacked_type = uint64_t;
using packed_type = uint64_2;
static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } // this replication makes code cleaner later.
};
template<> struct packing<float4> {
static __device__ inline constexpr int num() { return 4; }
};
template<> struct packing<int4> {
static __device__ inline constexpr int num() { return 4; }
};
#ifdef KITTENS_HOPPER
template<> struct packing<fp8e4m3> {
static __device__ inline constexpr int num() { return 1; }
using unpacked_type = fp8e4m3;
using packed_type = fp8e4m3_4;
};
template<> struct packing<fp8e4m3_4> {
static __device__ inline constexpr int num() { return 4; }
using unpacked_type = fp8e4m3;
using packed_type = fp8e4m3_4;
};
template<> struct packing<fp8e5m2> {
static __device__ inline constexpr int num() { return 1; }
using unpacked_type = fp8e5m2;
using packed_type = fp8e5m2_4;
};
template<> struct packing<fp8e5m2_4> {
static __device__ inline constexpr int num() { return 4; }
using unpacked_type = fp8e5m2;
using packed_type = fp8e5m2_4;
};
#ifdef KITTENS_BLACKWELL
template<> struct packing<fp8e8m0> {
static __device__ inline constexpr int num() { return 1; }
using unpacked_type = fp8e8m0;
using packed_type = fp8e8m0_4;
};
template<> struct packing<fp8e8m0_4> {
static __device__ inline constexpr int num() { return 4; }
using unpacked_type = fp8e8m0;
using packed_type = fp8e8m0_4;
};
#endif
#endif
/**
* @brief Provides templated functionality to convert between different types.
*
* @tparam T The target type for conversion.
* @tparam U The source type for conversion.
*/
template<typename T, typename U> struct convertor {
/**
* @brief Converts a value of type U to type T.
*
* @param u[in] The value of type U to convert.
* @return T The converted value of type T.
*/
static __host__ __device__ inline T convert(const U & u) {
return (T)u;
}
};
template<> struct convertor<float, bf16> {
static __host__ __device__ inline float convert(const bf16 & u) {
return __bfloat162float(u);
}
};
template<> struct convertor<bf16, float> {
static __host__ __device__ inline bf16 convert(const float & u) {
return __float2bfloat16_rn(u);
}
};
template<> struct convertor<float2, bf16_2> {
static __host__ __device__ inline float2 convert(const bf16_2 & u) {
return __bfloat1622float2(u);
}
};
template<> struct convertor<bf16_2, float2> {
static __host__ __device__ inline bf16_2 convert(const float2 & u) {
return __float22bfloat162_rn(u);
}
};
template<> struct convertor<float, half> {
static __host__ __device__ inline float convert(const half & u) {
return __half2float(u);
}
};
template<> struct convertor<half, float> {
static __host__ __device__ inline half convert(const float & u) {
return __float2half(u);
}
};
template<> struct convertor<float2, half_2> {
static __host__ __device__ inline float2 convert(const half_2 & u) {
return __half22float2(u);
}
};
template<> struct convertor<half_2, float2> {
static __host__ __device__ inline half_2 convert(const float2 & u) {
return __float22half2_rn(u);
}
};
template<> struct convertor<bf16, half> {
static __host__ __device__ inline bf16 convert(const half & u) {
return __float2bfloat16_rn(__half2float(u));
}
};
template<> struct convertor<half, bf16> {
static __host__ __device__ inline half convert(const bf16 & u) {
return __float2half(__bfloat162float(u));
}
};
template<> struct convertor<bf16_2, half_2> {
static __host__ __device__ inline bf16_2 convert(const half_2 & u) {
return __float22bfloat162_rn(__half22float2(u));
}
};
template<> struct convertor<half_2, bf16_2> {
static __host__ __device__ inline half_2 convert(const bf16_2 & u) {
return __float22half2_rn(__bfloat1622float2(u));
}
};
#ifdef KITTENS_HOPPER
// fp8e4m3
template<> struct convertor<fp8e4m3_4, float4> {
static __host__ __device__ inline fp8e4m3_4 convert(const float4& u) {
return __nv_fp8x4_e4m3(u);
}
};
template<> struct convertor<float4, fp8e4m3_4> {
static __host__ __device__ inline float4 convert(const fp8e4m3_4& u) {
__nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x4_e4m3*>(&u));
return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3]));
}
};
template<> struct convertor<fp8e4m3_2, float2> {
static __host__ __device__ inline fp8e4m3_2 convert(const float2& u) {
return __nv_fp8x2_e4m3(u);
}
};
template<> struct convertor<float2, fp8e4m3_2> {
static __host__ __device__ inline float2 convert(const fp8e4m3_2& u) {
__nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x2_e4m3*>(&u));
return make_float2(float(vals[0]), float(vals[1]));
}
};
template<> struct convertor<fp8e4m3, float> {
static __host__ __device__ inline fp8e4m3 convert(const float & u) {
return __nv_fp8_e4m3(u);
}
};
template<> struct convertor<float, fp8e4m3> {
static __host__ __device__ inline float convert(const fp8e4m3 & u) {
return float(u);
}
};
template<> struct convertor<bf16_2, fp8e4m3_4> {
static __host__ __device__ inline bf16_2 convert(const fp8e4m3_4 & u) {
float4 f4 = convertor<float4, fp8e4m3_4>::convert(u);
float2 f2 = make_float2(f4.x, f4.y);
return __float22bfloat162_rn(f2);
}
};
template<> struct convertor<fp8e4m3_4, bf16_2> {
static __host__ __device__ inline fp8e4m3_4 convert(const bf16_2 & u) {
float2 f2 = __bfloat1622float2(u);
float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f);
return __nv_fp8x4_e4m3(f4);
}
};
// fp8e5m2
template<> struct convertor<fp8e5m2_4, float4> {
static __host__ __device__ inline fp8e5m2_4 convert(const float4& u) {
return __nv_fp8x4_e5m2(u);
}
};
template<> struct convertor<float4, fp8e5m2_4> {
static __host__ __device__ inline float4 convert(const fp8e5m2_4& u) {
__nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x4_e5m2*>(&u));
return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3]));
}
};
template<> struct convertor<fp8e5m2_2, float2> {
static __host__ __device__ inline fp8e5m2_2 convert(const float2& u) {
return __nv_fp8x2_e5m2(u);
}
};
template<> struct convertor<float2, fp8e5m2_2> {
static __host__ __device__ inline float2 convert(const fp8e5m2_2& u) {
__nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x2_e5m2*>(&u));
return make_float2(float(vals[0]), float(vals[1]));
}
};
template<> struct convertor<fp8e5m2, float> {
static __host__ __device__ inline fp8e5m2 convert(const float & u) {
return __nv_fp8_e5m2(u);
}
};
template<> struct convertor<float, fp8e5m2> {
static __host__ __device__ inline float convert(const fp8e5m2 & u) {
return float(u);
}
};
template<> struct convertor<bf16_2, fp8e5m2_4> {
static __host__ __device__ inline bf16_2 convert(const fp8e5m2_4 & u) {
float4 f4 = convertor<float4, fp8e5m2_4>::convert(u);
float2 f2 = make_float2(f4.x, f4.y);
return __float22bfloat162_rn(f2);
}
};
template<> struct convertor<fp8e5m2_4, bf16_2> {
static __host__ __device__ inline fp8e5m2_4 convert(const bf16_2 & u) {
float2 f2 = __bfloat1622float2(u);
float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f);
return __nv_fp8x4_e5m2(f4);
}
};
#endif
}
}
@@ -1,11 +0,0 @@
/**
* @file
* @brief A collection of common resources on which ThunderKittens depends.
*/
#pragma once
#include "util.cuh"
#include "base_types.cuh"
#include "base_ops.cuh"
@@ -1,56 +0,0 @@
#pragma once
// Reset
#define TK_RESET "\033[0m"
// Foreground colors
#define TK_FG_BLACK "\033[30m"
#define TK_FG_RED "\033[31m"
#define TK_FG_GREEN "\033[32m"
#define TK_FG_YELLOW "\033[33m"
#define TK_FG_BLUE "\033[34m"
#define TK_FG_MAGENTA "\033[35m"
#define TK_FG_CYAN "\033[36m"
#define TK_FG_WHITE "\033[37m"
// Background colors
#define TK_BG_BLACK "\033[40m"
#define TK_BG_RED "\033[41m"
#define TK_BG_GREEN "\033[42m"
#define TK_BG_YELLOW "\033[43m"
#define TK_BG_BLUE "\033[44m"
#define TK_BG_MAGENTA "\033[45m"
#define TK_BG_CYAN "\033[46m"
#define TK_BG_WHITE "\033[47m"
// Bright foreground colors
#define TK_FG_BRIGHT_BLACK "\033[90m"
#define TK_FG_BRIGHT_RED "\033[91m"
#define TK_FG_BRIGHT_GREEN "\033[92m"
#define TK_FG_BRIGHT_YELLOW "\033[93m"
#define TK_FG_BRIGHT_BLUE "\033[94m"
#define TK_FG_BRIGHT_MAGENTA "\033[95m"
#define TK_FG_BRIGHT_CYAN "\033[96m"
#define TK_FG_BRIGHT_WHITE "\033[97m"
// Bright background colors
#define TK_BG_BRIGHT_BLACK "\033[100m"
#define TK_BG_BRIGHT_RED "\033[101m"
#define TK_BG_BRIGHT_GREEN "\033[102m"
#define TK_BG_BRIGHT_YELLOW "\033[103m"
#define TK_BG_BRIGHT_BLUE "\033[104m"
#define TK_BG_BRIGHT_MAGENTA "\033[105m"
#define TK_BG_BRIGHT_CYAN "\033[106m"
#define TK_BG_BRIGHT_WHITE "\033[107m"
// Text styles
#define TK_BOLD "\033[1m"
#define TK_DIM "\033[2m"
#define TK_ITALIC "\033[3m"
#define TK_UNDERLINE "\033[4m"
#define TK_BLINK "\033[5m"
#define TK_REVERSE "\033[7m"
#define TK_HIDDEN "\033[8m"
// Macro to combine styles
#define TK_STYLE(...) "\033[" #__VA_ARGS__ "m"
-314
View File
@@ -1,314 +0,0 @@
/**
* @file
* @brief General utilities for ThunderKittens.
*/
#pragma once
#include <stdint.h>
#include <type_traits>
#include <concepts>
#include <memory>
// CUDA driver API
#define CUCHECK(cmd) do { \
CUresult err = cmd; \
if (err != CUDA_SUCCESS) { \
const char *errStr; \
cuGetErrorString(err, &errStr); \
fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n", \
__FILE__, __LINE__, errStr); \
exit(EXIT_FAILURE); \
} \
} while(0)
// CUDA runtime API
#define CUDACHECK(cmd) do { \
cudaError_t err = cmd; \
if (err != cudaSuccess) { \
fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n", \
__FILE__, __LINE__, cudaGetErrorString(err)); \
exit(EXIT_FAILURE); \
} \
} while(0)
/**
* @namespace kittens
*
* @brief The main namespace of ThunderKittens.
*/
namespace kittens {
/* ---------- GENERAL CONSTANTS FOR KITTENS ---------- */
/**
* @brief Tile dimension constant.
*/
template<typename T> constexpr int TILE_COL_DIM = sizeof(T) == 1 ? 32 : 16;
template<typename T> constexpr int TILE_ROW_DIM = 16;
/**
* @brief Tile num elements constant calculated as TILE_DIM squared.
*/
template<typename T> constexpr int TILE_ELEMENTS{TILE_COL_DIM<T>*TILE_ROW_DIM<T>};
/**
* @brief Constant representing number of threads in a warp.
*/
constexpr int WARP_THREADS{32};
/**
* @brief Constant representing number of threads in a warpgroup of four warps.
*/
constexpr int WARPGROUP_THREADS{128};
/**
* @brief Constant representing number of warps in a warpgroup of four warps.
*/
constexpr int WARPGROUP_WARPS{4};
/**
* @brief Get the warp ID of the current thread.
* @return The warp ID.
*/
__device__ static __forceinline__ int warpid() {
// uint32_t wid;
// asm volatile("mov.u32 %0, %warpid;" : "=r"(wid));
// return wid;
return threadIdx.x >> 5;
}
/**
* @brief Get the warpgroup ID of the current thread.
* @return The warpgroup ID.
*/
__device__ static __forceinline__ int warpgroupid() { return warpid() >> 2; }
/**
* @brief Get the lane ID of the current thread within its warp.
* @return The lane ID.
*/
__device__ static __forceinline__ int laneid() {
// uint32_t lid;
// asm volatile("mov.u32 %0, %laneid;" : "=r"(lid));
// return lid;
return threadIdx.x & 31;
}
#if defined(KITTENS_HOPPER)
constexpr int MAX_SHARED_MEMORY = 227000;
#elif defined(KITTENS_A100)
constexpr int MAX_SHARED_MEMORY = 164000;
#elif defined(KITTENS_4090)
constexpr int MAX_SHARED_MEMORY = 100000;
#endif
struct transpose {
static constexpr int N = 0; // not transposed
static constexpr int T = 1; // transposed
};
struct axis {
static constexpr int ROW = 0; // row axis of a tile
static constexpr int COL = 1; // column axis of a tile
};
/* ---------- TYPE HELPERS ---------- */
/**
* @namespace ducks
*
* @brief ThunderKittens' namespace for template metaprogramming..
*
* This includes primarily dummy types and concept wrappers, along
* with a few additional utilities.
*/
namespace ducks {
/**
* @brief A type representing an empty default for a template.
*/
struct default_type {};
// This macro can't be done as a template, so it doesn't really have a location in kittens.
#define typeof(A) typename std::remove_const<typename std::remove_reference<decltype(A)>::type>::type
}
/* ---------- SHUFFLE UTILS ---------- */
/**
* @brief Mask constant for all active threads in a warp.
*/
static constexpr uint32_t MASK_ALL = 0xFFFFFFFF;
/**
* @brief Perform a shuffle down operation on a packed type synchronously across a warp.
* @tparam T The type of the value to be shuffled.
* @param mask[in] The mask of active threads.
* @param f[in] The value to be shuffled.
* @param delta[in] The number of positions to shuffle down.
* @return The result of the shuffle operation.
*/
template<typename T>
__device__ static inline T packed_shfl_down_sync(uint32_t mask, const T &f, int delta) {
return __shfl_down_sync(mask, f, delta);
}
template<>
__device__ inline float2 packed_shfl_down_sync<float2>(uint32_t mask, const float2 &f, int delta) {
float2 r;
r.x = __shfl_down_sync(mask, f.x, delta);
r.y = __shfl_down_sync(mask, f.y, delta);
return r;
}
/**
* @brief Perform a packed shuffle operation synchronously across a warp.
* @tparam T The type of the value to be shuffled.
* @param mask[in] The mask of active threads.
* @param f[in] The value to be shuffled.
* @param src[in] The source lane from which to shuffle.
* @return The result of the shuffle operation.
*/
template<typename T>
__device__ static inline T packed_shfl_sync(uint32_t mask, const T &f, int src) {
return __shfl_sync(mask, f, src);
}
template<>
__device__ inline float2 packed_shfl_sync<float2>(uint32_t mask, const float2 &f, int src) {
float2 r;
r.x = __shfl_sync(mask, f.x, src);
r.y = __shfl_sync(mask, f.y, src);
return r;
}
/* ---------- SHARED MEMORY UTILS ---------- */
// namespace ducks {
// namespace sb {
// struct identifier {};
// }
// }
// template<typename Args...>
// struct sb {
// using identifier = ducks::sb::identifier;
// Args... args;
// };
// namespace ducks {
// namespace sb {
// template<typename T> concept all = requires {
// typename T::identifier;
// } && std::is_same_v<T::identifier, identifier>;
// }
// }
// Joyously stolen from https://github.com/NVIDIA/cutlass/blob/5c447dd84f8ae0e1d48ff9a2eae26ce8c4958101/include/cute/container/alignment.hpp#L51
#if defined(__CUDACC__)
#define KITTENS_ALIGN_AS(n) __align__(n)
#else
#define KITTENS_ALIGN_AS(n) alignas(n)
#endif
#ifdef KITTENS_HOPPER
#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(128)
#else
#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(16)
#endif
/**
* @brief Dummy structure for alignment purposes. Needed for WGMMA and TMA calls.
*/
struct KITTENS_DEFAULT_ALIGN alignment_dummy { int dummy; };
/**
* @brief Very simple allocator for dynamic shared memory. Advances pointer and tracks alignments.
* @tparam default_alignment The default alignment this allocator will enforce. If <=0 (default -1) it will not align.
*/
#ifdef KITTENS_HOPPER
template<int default_alignment=1024>
#else
template<int default_alignment=16>
#endif
struct shared_allocator {
int *ptr;
private:
// Recursive template to generate N-dimensional array type
template<typename A, size_t... dims>
struct variadic_array;
template<typename A, size_t first_dim, size_t... rest_dims>
struct variadic_array<A, first_dim, rest_dims...> {
using type = typename variadic_array<A, rest_dims...>::type[first_dim];
};
template<typename A>
struct variadic_array<A> {
using type = A;
};
template<typename A, size_t... dims>
using variadic_array_t = typename variadic_array<A, dims...>::type;
template<int alignment>
__device__ inline void align_ptr() {
if constexpr (alignment > 0) {
uint64_t p = reinterpret_cast<uint64_t>(ptr);
if(p % alignment != 0) {
ptr = (int*)(p + (alignment-(p%alignment)));
}
}
}
public:
/**
* @brief Construct a new shared allocator using a pointer to extern shared memory.
* @param[in] _ptr Pointer to the start of the extern shared memory.
*/
__device__ shared_allocator(int *_ptr): ptr(_ptr) {}
/**
* @brief Allocate shared memory for a single instance or N-dimensional array of type A.
* @tparam A The type of the object to allocate.
* @tparam dims... A list of dimensions for the N-dimensional array.
* @return Reference to the allocated object.
*/
template<typename A, size_t... dims>
__device__ inline variadic_array_t<A, dims...>& allocate() {
// static_assert(sizeof(A) % default_alignment == 0, "Type is not aligned properly for array allocation");
align_ptr<default_alignment>();
using at = variadic_array_t<A, dims...>;
at*p = reinterpret_cast<at*>(ptr);
ptr += sizeof(at)/sizeof(int);
return *p;
}
/**
* @brief Allocate shared memory for a single instance or N-dimensional array of type A.
* @tparam alignment An alignment to enforce for this particular object.
* @tparam A The type of the object to allocate.
* @tparam dims... A list of dimensions for the N-dimensional array.
* @return Reference to the allocated object.
*/
template<int alignment, typename A, size_t... dims>
__device__ inline variadic_array_t<A, dims...>& allocate() {
// static_assert(sizeof(A) % alignment == 0, "Type is not aligned properly for array allocation");
align_ptr<alignment>();
using at = variadic_array_t<A, dims...>;
at*p = reinterpret_cast<at*>(ptr);
ptr += sizeof(at)/sizeof(int);
return *p;
}
};
#if (defined(KITTENS_HOPPER) || defined(KITTENS_BLACKWELL))
/**
* @brief A wrapper for an allocator that enforces sufficient alignment to be used for TMA loads and stores.
*/
using tma_allocator = shared_allocator<1024>;
using tma_swizzle_allocator = tma_allocator; // swizzled TMA modes require up to 1024 byte alignments :/
/* Get CTA ID within a cluster */
__device__ static inline int3 clusterIdx() {
int3 cluster_idx;
asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(cluster_idx.x));
asm volatile("mov.u32 %0, %clusterid.y;\n" : "=r"(cluster_idx.y));
asm volatile("mov.u32 %0, %clusterid.z;\n" : "=r"(cluster_idx.z));
return cluster_idx;
}
__device__ static inline int cluster_ctarank() {
uint32_t ctarank;
asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(ctarank));
return ctarank;
}
#endif
} // namespace kittens
-12
View File
@@ -1,12 +0,0 @@
/**
* @file
* @brief The master header file of ThunderKittens. This file includes everything you need!
*/
#pragma once
#include "common/common.cuh"
#include "types/types.cuh"
#include "ops/ops.cuh"
#include "pyutils/util.cuh"
// #include "pyutils/pyutils.cuh" // for simple binding without including torch
@@ -1,51 +0,0 @@
/**
* @file
* @brief An aggregate header of all device (multi-GPU) operations defined by ThunderKittens
*/
#pragma once
#include "../../types/types.cuh"
namespace kittens {
template<int _NUM_DEVICES>
struct device {
static_assert(_NUM_DEVICES >= 0 && _NUM_DEVICES <= 72, "Invalid number of devices");
static constexpr int NUM_DEVICES = _NUM_DEVICES;
#ifdef KITTENS_HOPPER
using barrier_t = pgl<gl<int, 1, 1, 1, -1>, NUM_DEVICES, true>;
/**
* @brief Multi-GPU synchronization barrier for coordinated kernel exit
*
* Performs a synchronization across all devices to ensure all GPUs complete
* their work before any kernel exits. Does not synchronize intra-node threads
* or threadblocks.
*
* @param barrier Pre-allocated barrier structure, must be initialized to 0
* @param dev_idx Current device index (0 to NUM_DEVICES - 1)
* @param id Synchronization point identifier (default: 0). 0 is fine for most cases
*
*/
__device__ static inline void sync_on_exit(const barrier_t &barrier, const int dev_idx, const int id = 0) {
if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 &&
threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) {
cuda::atomic_ref<int, cuda::thread_scope_system> barrier_uc(barrier[dev_idx][{id}]);
// Inter-note check-in
multimem<int>::red<reduce_op::ADD>(barrier.mc_ptr_at({id}), 1);
asm volatile ("{fence.proxy.alias;}" ::: "memory");
while (barrier_uc.load(cuda::memory_order_acquire) < NUM_DEVICES);
barrier_uc.fetch_sub(NUM_DEVICES, cuda::memory_order_release);
}
}
#endif
};
} // namespace kittens
@@ -1,96 +0,0 @@
/**
* @file
* @brief An aggregate header of all group (multi-warp) operations defined by ThunderKittens
*/
#pragma once
#include <cuda/pipeline>
#include "../../common/common.cuh"
#include "../../types/types.cuh"
#include "../thread/thread.cuh" // several group memory ops rely on underlying warp-scope ops
#define KITTENS_CHECK_WARP static_assert(GROUP_WARPS==1, "Warp (GROUP_WARPS=1) function called from a non-warp group.");
// A "warpgroup" is a special group of 4 consecutive warps defined by NVIDIA for certain SM_90+ operations.
#define KITTENS_CHECK_WARPGROUP static_assert(GROUP_WARPS==4, "Warpgroup (GROUP_WARPS=4) function called from a non-warpgroup group.");
// WGMMA relies on some template structures that cannot be specialized within the group struct, so we declare them in advance.
#ifdef KITTENS_HOPPER
#include "mma/warpgroup/base/base.cuh"
#endif
namespace kittens {
/*
This is meant to be used with a `using group_N = kittens::group<NUM_WORKERS>;` at the start of every kernel.
*/
template<int _GROUP_WARPS>
struct group {
static constexpr int GROUP_WARPS = _GROUP_WARPS; // This alias produces nice parallelism.
static constexpr int GROUP_THREADS = GROUP_WARPS * kittens::WARP_THREADS; // This alias produces nice parallelism.
__device__ static inline int laneid() { return threadIdx.x % GROUP_THREADS; }
__device__ static inline int warpid() { return laneid() / kittens::WARP_THREADS; }
__device__ static inline int groupid() { return threadIdx.x / GROUP_THREADS; }
__device__ static inline void sync(int id) {
asm volatile("bar.sync %0, %1;\n" :: "r"(id), "n"(GROUP_THREADS));
}
template<uint32_t MASK=0xFFFFFFFF> __device__ static inline void sync() {
static_assert(GROUP_WARPS==1, "barrier-less sync() can only be called by a single warp!");
asm volatile("bar.warp.sync %0;\n" :: "n"(MASK));
}
__device__ static inline void arrive(int id) {
asm volatile("bar.arrive %0, %1;\n" :: "r"(id), "n"(GROUP_THREADS));
}
#include "memory/memory.cuh"
#include "shared/shared.cuh"
#include "register/register.cuh"
#ifdef KITTENS_HOPPER
#include "mma/mma.cuh"
template<int n_reg> __device__ static inline void increase_registers() {
static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8");
asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" :: "n"(n_reg));
}
template<int n_reg> __device__ static inline void decrease_registers() {
static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8");
asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" :: "n"(n_reg));
}
__device__ static inline void producer_registers() { decrease_registers<24>(); }
template<int NCWG> __device__ static inline void consumer_registers() { increase_registers<480/NCWG - 8*(NCWG>3) - 224*(NCWG==1)>(); }
#endif
};
namespace everyone {
// Block-level synchronization
__device__ static inline void sync(int id) {
asm volatile("bar.sync %0;\n" :: "r"(id));
}
// Cluster-level synchronization functions
namespace tma {
namespace cluster {
__device__ static inline void arrive_aligned() { // All threads in the cluster must call this
asm volatile ("barrier.cluster.arrive.release.aligned;\n");
}
__device__ static inline void wait_aligned() {
asm volatile ("barrier.cluster.wait.acquire.aligned;\n");
}
__device__ static inline void sync() {
arrive_aligned();
wait_aligned();
}
}
}
};
using warp = group<1>; // scope used by most pre-Hopper GPUs, and also for most register operations.
using warpgroup = group<4>; // special scope commonly used by Hopper and later.
}
@@ -1,21 +0,0 @@
/**
* @file
* @brief An aggregate header of colaborative group memory movement operations
*/
#include "util/util.cuh"
#include "tile/tile.cuh"
#include "vec/vec.cuh"
#ifdef KITTENS_HOPPER
struct tma {
#include "util/tma.cuh"
#include "tile/tma.cuh"
#include "vec/tma.cuh"
struct cluster {
#include "util/tma_cluster.cuh"
#include "tile/tma_cluster.cuh"
#include "vec/tma_cluster.cuh"
};
};
#endif
@@ -1,42 +0,0 @@
/**
* @file
* @brief Functions for a group to collaboratively transfer data directly between global memory and registers and back.
*/
/**
* @brief Collaboratively loads data from a source array into register tiles.
*
* @tparam RT The register tile type.
* @tparam U The data type of the source array.
* @param dst[out] The destination tile to load data into.
* @param src[in] The source array to load data from.
* @param row_stride[in] The stride in elements between rows in the source array.
*/
template<int axis, ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
__device__ inline static void load(CRT &dst, const CGL &src, const COORD &idx) {
load<axis, CRT::component, CGL::component, COORD>(dst.real, src.real, idx);
load<axis, CRT::component, CGL::component, COORD>(dst.imag, src.imag, idx);
}
template<ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
__device__ inline static void load(CRT &dst, const CGL &src, const COORD &idx) {
load<2, CRT, CGL>(dst, src, idx);
}
/**
* @brief Collaboratively stores data from register tiles to a destination array in global memory.
*
* @tparam RT The register tile type.
* @tparam U The data type of the destination array.
* @param[out] dst The destination array in global memory to store data into.
* @param[in] src The source register tile to store data from.
* @param row_stride[in] The stride in elements between rows in the destination array.
*/
template<int axis, ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
__device__ inline static void store(CGL &dst, const CRT &src, const COORD &idx) {
store<axis, typename CRT::component, typename CGL::component>(dst.real, src.real, idx);
store<axis, typename CRT::component, typename CGL::component>(dst.imag, src.imag, idx);
}
template<ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
__device__ inline static void store(CGL &dst, const CRT &src, const COORD &idx) {
store<2, CRT, CGL>(dst, src, idx);
}
@@ -1,37 +0,0 @@
/**
* @file
* @brief Group (collaborative warp) ops for loading shared tiles from and storing to global memory.
*/
template<int axis, bool assume_aligned, ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
__device__ static inline void load(CST &dst, const CGL &src, const COORD &idx) {
load<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
load<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
}
template<ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
__device__ static inline void load(CST &dst, const CGL &src, const COORD &idx) {
load<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
load<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
}
template<int axis, bool assume_aligned, ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
__device__ static inline void store(CGL &dst, const CST &src, const COORD &idx) {
store<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
store<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
}
template<ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
__device__ static inline void store(CGL &dst, const CST &src, const COORD &idx) {
store<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
store<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
}
template<int axis, bool assume_aligned, ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
__device__ static inline void load_async(CST &dst, const CGL &src, const COORD &idx) {
load_async<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
load_async<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
}
template<ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
__device__ static inline void load_async(CST &dst, const CGL &src, const COORD &idx) {
load_async<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
load_async<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
}
@@ -1,34 +0,0 @@
/**
* @file
* @brief Functions for a warpgroup to collaboratively transfer data directly between shared memory and registers and back.
*/
/**
* @brief Collaboratively load data from a shared tile into register tiles split across a warpgroup.
*
* @tparam RT The register tile type
* @tparam ST The shared tile type
* @param dst[out] The destination register tile.
* @param src[in] The source shared tile.
*/
template<ducks::crt::all RT, ducks::cst::all ST>
__device__ inline static void load(RT &dst, const ST &src) {
load(dst.real, src.real);
load(dst.imag, src.imag);
}
/**
* @brief Collaboratively store data into a shared tile from register tiles split across a warpgroup.
*
* @tparam RT The register tile type
* @tparam ST The shared tile type
* @param dst[out] The destination shared tile.
* @param src[in] The source register tile.
*/
template<ducks::cst::all ST, ducks::crt::all RT>
__device__ inline static void store(ST &dst, const RT &src) {
store(dst.real, src.real);
store(dst.imag, src.imag);
}
@@ -1,207 +0,0 @@
/**
* @file
* @brief Functions for a group to collaboratively transfer data directly between global memory and registers and back.
*/
/**
* @brief Collaboratively loads data from a source array into row-major layout tiles.
*
* @tparam RT The row-major layout tile type.
* @tparam U The data type of the source array.
* @param dst[out] The destination tile to load data into.
* @param src[in] The source array to load data from.
* @param row_stride[in] The stride in elements between rows in the source array.
*/
template<int axis, ducks::rt::row_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) {
using T2 = RT::dtype;
using U = typename GL::dtype;
#ifdef KITTENS_HOPPER
static_assert(!std::is_same_v<T2, fp8e4m3_4> && !std::is_same_v<T2, fp8e5m2_4>, "Unsupported type for load/store");
#endif
U *src_ptr = (U*)&src[(idx.template unit_coord<axis, 3>())];
const int row_stride = src.template stride<axis>();
using U2 = base_types::packing<U>::packed_type;
int warp_laneid = threadIdx.x % WARP_THREADS;
int local_warpid;
if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
else local_warpid = warpid();
const int row_offset = dst.rows*local_warpid;
#pragma unroll
for(int i = 0; i < dst.height; i++) {
int row = row_offset + i*dst.tile_size_row + (warp_laneid / 4);
#pragma unroll
for(int j = 0; j < dst.width; j++) {
int col = j*dst.tile_size_col + 2*(warp_laneid % 4);
dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+0)*row_stride + (col+0)]));
dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+0)*row_stride + (col+8)]));
}
#pragma unroll
for(int j = 0; j < dst.width; j++) {
int col = j*dst.tile_size_col + 2*(warp_laneid % 4);
dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+8)*row_stride + (col+0)]));
dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+8)*row_stride + (col+8)]));
}
}
}
/**
* @brief Collaboratively loads data from a source array into column-major layout tiles.
*
* @tparam RT The column-major layout tile type.
* @tparam U The data type of the source array.
* @param dst[out] The destination tile to load data into.
* @param src[in] The source array to load data from.
* @param row_stride[in] The stride in elements between rows in the source array.
*/
template<int axis, ducks::rt::col_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) {
using T = typename RT::T;
using U = typename GL::dtype;
#ifdef KITTENS_HOPPER
static_assert(!std::is_same_v<T, fp8e4m3> && !std::is_same_v<T, fp8e5m2>, "Unsupported type for load/store");
#endif
U *src_ptr = (U*)&src[(idx.template unit_coord<axis, 3>())];
const int row_stride = src.template stride<axis>();
int warp_laneid = threadIdx.x % WARP_THREADS;
int local_warpid;
if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
else local_warpid = warpid();
const int row_offset = dst.rows*local_warpid;
#pragma unroll
for(int i = 0; i < dst.height; i++) {
int row = row_offset + i*dst.tile_size_row + 2*(warp_laneid % 4);
#pragma unroll
for(int j = 0; j < dst.width; j++) {
int col = j*dst.tile_size_col + (warp_laneid / 4);
dst.tiles[i][j].data[0].x = base_types::convertor<T, U>::convert(src_ptr[(row+0)*row_stride + (col+0)]);
dst.tiles[i][j].data[1].x = base_types::convertor<T, U>::convert(src_ptr[(row+0)*row_stride + (col+8)]);
}
#pragma unroll
for(int j = 0; j < dst.width; j++) {
int col = j*dst.tile_size_col + (warp_laneid / 4);
dst.tiles[i][j].data[0].y = base_types::convertor<T, U>::convert(src_ptr[(row+1)*row_stride + (col+0)]);
dst.tiles[i][j].data[1].y = base_types::convertor<T, U>::convert(src_ptr[(row+1)*row_stride + (col+8)]);
}
#pragma unroll
for(int j = 0; j < dst.width; j++) {
int col = j*dst.tile_size_col + (warp_laneid / 4);
dst.tiles[i][j].data[2].x = base_types::convertor<T, U>::convert(src_ptr[(row+8)*row_stride + (col+0)]);
dst.tiles[i][j].data[3].x = base_types::convertor<T, U>::convert(src_ptr[(row+8)*row_stride + (col+8)]);
}
#pragma unroll
for(int j = 0; j < dst.width; j++) {
int col = j*dst.tile_size_col + (warp_laneid / 4);
dst.tiles[i][j].data[2].y = base_types::convertor<T, U>::convert(src_ptr[(row+9)*row_stride + (col+0)]);
dst.tiles[i][j].data[3].y = base_types::convertor<T, U>::convert(src_ptr[(row+9)*row_stride + (col+8)]);
}
}
}
template<ducks::rt::all RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) {
load<2>(dst, src, idx);
}
/**
* @brief Collaboratively stores data from register tiles to a destination array in global memory with a row-major layout.
*
* @tparam RT The register tile type with a row-major layout.
* @tparam U The data type of the destination array.
* @param[out] dst The destination array in global memory to store data into.
* @param[in] src The source register tile to store data from.
* @param row_stride[in] The stride in elements between rows in the destination array.
*/
template<int axis, ducks::rt::row_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) {
using T2 = RT::dtype;
using U = typename GL::dtype;
#ifdef KITTENS_HOPPER
static_assert(!std::is_same_v<T2, fp8e4m3_4> && !std::is_same_v<T2, fp8e5m2_4>, "Unsupported type for load/store");
#endif
U *dst_ptr = (U*)&dst[(idx.template unit_coord<axis, 3>())];
const int row_stride = dst.template stride<axis>();
using U2 = base_types::packing<U>::packed_type;
int warp_laneid = threadIdx.x % WARP_THREADS;
int local_warpid;
if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
else local_warpid = warpid();
const int row_offset = src.rows*local_warpid;
#pragma unroll
for(int i = 0; i < src.height; i++) {
int row = row_offset + i*src.tile_size_row + (warp_laneid / 4);
#pragma unroll
for(int j = 0; j < src.width; j++) {
int col = j*src.tile_size_col + 2*(warp_laneid % 4);
*(U2*)(&dst_ptr[(row+0)*row_stride + (col+0)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
*(U2*)(&dst_ptr[(row+0)*row_stride + (col+8)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
}
#pragma unroll
for(int j = 0; j < src.width; j++) {
int col = j*src.tile_size_col + 2*(warp_laneid % 4);
*(U2*)(&dst_ptr[(row+8)*row_stride + (col+0)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
*(U2*)(&dst_ptr[(row+8)*row_stride + (col+8)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
}
}
}
/**
* @brief Collaboratively stores data from register tiles to a destination array in global memory with a column-major layout.
*
* @tparam RT The register tile type with a column-major layout.
* @tparam U The data type of the destination array.
* @param[out] dst The destination array in global memory to store data into.
* @param[in] src The source register tile to store data from.
* @param row_stride[in] The stride in elements between rows in the destination array.
*/
template<int axis, ducks::rt::col_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) {
using T = base_types::packing<typename RT::dtype>::unpacked_type;
using U = typename GL::dtype;
#ifdef KITTENS_HOPPER
static_assert(!std::is_same_v<T, fp8e4m3_4> && !std::is_same_v<T, fp8e5m2_4>, "Unsupported type for load/store");
#endif
U *dst_ptr = (U*)&dst[(idx.template unit_coord<axis, 3>())];
const int row_stride = dst.template stride<axis>();
int warp_laneid = threadIdx.x % WARP_THREADS;
int local_warpid;
if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
else local_warpid = warpid();
const int row_offset = src.rows*local_warpid;
#pragma unroll
for(int i = 0; i < src.height; i++) {
int row = row_offset + i*src.tile_size_row + 2*(warp_laneid % 4);
#pragma unroll
for(int j = 0; j < src.width; j++) {
int col = j*src.tile_size_col + (warp_laneid / 4);
dst_ptr[(row+0)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[0].x);
dst_ptr[(row+0)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[1].x);
}
#pragma unroll
for(int j = 0; j < src.width; j++) {
int col = j*src.tile_size_col + (warp_laneid / 4);
dst_ptr[(row+1)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[0].y);
dst_ptr[(row+1)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[1].y);
}
#pragma unroll
for(int j = 0; j < src.width; j++) {
int col = j*src.tile_size_col + (warp_laneid / 4);
dst_ptr[(row+8)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[2].x);
dst_ptr[(row+8)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[3].x);
}
#pragma unroll
for(int j = 0; j < src.width; j++) {
int col = j*src.tile_size_col + (warp_laneid / 4);
dst_ptr[(row+9)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[2].y);
dst_ptr[(row+9)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[3].y);
}
}
}
template<ducks::rt::all RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) {
store<2>(dst, src, idx);
}
@@ -1,168 +0,0 @@
/**
* @file
* @brief Group (collaborative warp) ops for loading shared tiles from and storing to global memory.
*/
/**
* @brief Loads data from global memory into a shared memory tile.
*
* @tparam ST The type of the shared tile.
* @param[out] dst The destination shared memory tile.
* @param[in] src The source global memory array.
* @param[in] idx The coordinate of the tile in the global memory array.
*/
template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) {
using T = typename ST::dtype;
const int row_stride = src.template stride<axis>();
// we can handle this many rows each time we run a memcpy_async
constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype);
constexpr int memcpy_per_row = dst.cols / elem_per_memcpy;
constexpr int total_calls = (dst.height*dst.width * kittens::TILE_ROW_DIM<T>*kittens::TILE_COL_DIM<T> + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); // round up
constexpr int total_rows = dst.height*dst.width;
coord<> unit_coord = idx.template unit_coord<axis, 3>();
typename GL::dtype *src_ptr = (typename GL::dtype*)&src[unit_coord];
uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
int laneid = threadIdx.x % GROUP_THREADS;
#pragma unroll
for(int i = 0; i < total_calls; i++) {
int load_idx = i * GROUP_THREADS + laneid;
int row = load_idx / memcpy_per_row;
int col = (load_idx*elem_per_memcpy) % dst.cols;
if constexpr (assume_aligned) {
float4 tmp;
move<float4>::ldg(tmp, (float4*)&src_ptr[row*row_stride + col]);
move<float4>::sts(dst.idx(dst_ptr, {row, col}), tmp);
}
else {
if (row + unit_coord.template dim<axis>() < src.template shape<axis>()) {
float4 tmp;
move<float4>::ldg(tmp, (float4*)&src_ptr[row*row_stride + col]);
move<float4>::sts(dst.idx(dst_ptr, {row, col}), tmp);
}
else {
float4 zeros = {0.f,0.f,0.f,0.f};
move<float4>::sts(dst.idx(dst_ptr, {row, col}), zeros); // use the default value
}
}
}
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) {
load<2, false, ST, GL, COORD>(dst, src, idx);
}
/**
* @brief Stores data from a shared memory tile into global memory.
*
* @tparam ST The type of the shared tile.
* @param[out] dst The destination global memory array.
* @param[in] src The source shared memory tile.
* @param row_stride[in] The stride between rows in the destination array.
*/
template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) {
using T = typename ST::dtype;
const int row_stride = dst.template stride<axis>();
// we can handle this many rows each time we run a memcpy_async
constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype);
constexpr int memcpy_per_row = src.cols / elem_per_memcpy;
constexpr int total_calls = (src.height*src.width * kittens::TILE_ROW_DIM<T>*kittens::TILE_COL_DIM<T> + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); // round up
coord<> unit_coord = idx.template unit_coord<axis, 3>();
typename GL::dtype *dst_ptr = (typename GL::dtype*)&dst[unit_coord];
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
int laneid = threadIdx.x % GROUP_THREADS;
#pragma unroll
for(int i = 0; i < total_calls; i++) {
int load_idx = i * GROUP_THREADS + laneid;
int row = load_idx / memcpy_per_row;
int col = (load_idx*elem_per_memcpy) % src.cols;
if constexpr (assume_aligned) {
float4 tmp;
move<float4>::lds(tmp, src.idx(src_ptr, {row, col}));
move<float4>::stg((float4*)&dst_ptr[row*row_stride + col], tmp);
}
else {
if (row + unit_coord.template dim<axis>() < dst.template shape<axis>()) {
float4 tmp;
move<float4>::lds(tmp, src.idx(src_ptr, {row, col}));
move<float4>::stg((float4*)&dst_ptr[row*row_stride + col], tmp);
}
}
}
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) {
store<2, false, ST, GL, COORD>(dst, src, idx);
}
/**
* @brief Asynchronously loads data from global memory into a shared memory tile.
*
* @tparam ST The type of the shared tile.
* @param[out] dst The destination shared memory tile.
* @param[in] src The source global memory array.
*
* @note This function expects 16-byte alignments. Otherwise, behavior is undefined.
*/
template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx) {
using T = typename ST::dtype;
const int row_stride = src.template stride<axis>();
// we can handle this many rows each time we run a memcpy_async
constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype);
constexpr int memcpy_per_row = dst.cols / elem_per_memcpy;
constexpr int total_calls = (dst.height*dst.width * kittens::TILE_ROW_DIM<T>*kittens::TILE_COL_DIM<T> + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); // round up
coord<> unit_coord = idx.template unit_coord<axis, 3>();
typename GL::dtype *src_ptr = (typename GL::dtype*)&src[unit_coord];
uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
int laneid = threadIdx.x % GROUP_THREADS;
#pragma unroll
for(int i = 0; i < total_calls; i++) {
int load_idx = i * GROUP_THREADS + laneid;
int row = load_idx / memcpy_per_row;
int col = (load_idx*elem_per_memcpy) % dst.cols;
if constexpr (assume_aligned) {
asm volatile(
"cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n"
:: "r"(dst.idx(dst_ptr, {row, col})), "l"(&src_ptr[row*row_stride + col])
: "memory"
);
}
else {
if (row + unit_coord.template dim<axis>() < src.template shape<axis>()) {
asm volatile(
"cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n"
:: "r"(dst.idx(dst_ptr, {row, col})), "l"(&src_ptr[row*row_stride + col])
: "memory"
);
}
else {
// printf("thread %d skipping async load on row %d, col %d\n", threadIdx.x, row + unit_coord.template dim<axis>(), col);
float4 zeros = {0.f,0.f,0.f,0.f};
move<float4>::sts(dst.idx(dst_ptr, {row, col}), zeros); // use the default value
}
}
}
asm volatile("cp.async.commit_group;\n" ::: "memory");
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx) {
load_async<2, false, ST, GL, COORD>(dst, src, idx);
}
@@ -1,323 +0,0 @@
/**
* @file
* @brief Functions for a warpgroup to collaboratively transfer data directly between shared memory and registers and back.
*/
/**
* @brief Collaboratively load data from a shared tile into register tiles split across a warpgroup.
*
* @tparam RT The register tile type
* @tparam ST The shared tile type
* @param dst[out] The destination register tile.
* @param src[in] The source shared tile.
*/
template<ducks::rt::all RT, ducks::st::all ST>
__device__ inline static void load(RT &dst, const ST &src) {
constexpr int height = ST::height;
constexpr int warp_height = RT::height;
static_assert(height%GROUP_WARPS == 0, "Group load / store requires tile height to be a multiple of GROUP_WARPS.");
static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
int local_warpid;
if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
else local_warpid = warpid();
using T2 = RT::dtype;
using U = ST::dtype;
using T = base_types::packing<T2>::unpacked_type;
using U2 = base_types::packing<U>::packed_type;
int warp_laneid = ::kittens::laneid();
// convert to shared state space
uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
#pragma unroll
for(int i = 0; i < dst.height; i++) {
#pragma unroll
for(int j = 0; j < dst.width; j++) {
if constexpr (sizeof(typename ST::dtype) == 2) {
// handle the row-major layout for 16-bit types
U2 tmp[4];
int row = (local_warpid*warp_height + i)*dst.tile_size_row + (warp_laneid % 16);
int col = j*dst.tile_size_col + (warp_laneid / 16) * 8;
if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
move<U2>::ldsm4(tmp[0], tmp[1], tmp[2], tmp[3], src.idx(shared_addr, {row, col}));
}
else {
move<U2>::ldsm4t(tmp[0], tmp[2], tmp[1], tmp[3], src.idx(shared_addr, {row, col}));
}
dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
}
else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 1) {
// handle the row-major layout for 8-bit types
int warp_group_16 = (warp_laneid / 16); // divide each warp into two groups of 16 threads
int lane_in_16 = warp_laneid % 16; // position in group of 16 threads
int row = (local_warpid*warp_height + i)*dst.tile_size_row + (lane_in_16 % 16); // find base row for warp in warpgroup and then distribute the 16 threads in the warp across the rows
int col = j*dst.tile_size_col + warp_group_16 * 16; // find base column and then *16 for second half of the warp
U2 tmp[4];
if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
move<U2>::ldsm4(tmp[0], tmp[1], tmp[2], tmp[3], src.idx(shared_addr, {row, col}));
}
else {
move<U2>::ldsm4t(tmp[0], tmp[2], tmp[1], tmp[3], src.idx(shared_addr, {row, col}));
}
dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
}
else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 4) {
// handle the row-major layout for 32-bit types
int row = (local_warpid*warp_height + i)*dst.tile_size_row + (warp_laneid / 4);
int col = j*dst.tile_size_col + 2*(warp_laneid % 4);
if constexpr (ST::rows != ST::underlying_rows || ST::cols != ST::underlying_cols) { // subtile case
row += src.row_offset;
col += src.col_offset;
}
int blit = sizeof(typename ST::dtype) * ((warp_laneid%4) / 2);
U2 tmp[4];
static constexpr int swizzle_repeat = ST::swizzle_bytes * 8;
static constexpr int subtile_cols = ST::swizzle_bytes / sizeof(U);
const int outer_idx = col/subtile_cols;
const uint32_t addr_1 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+0)*subtile_cols + col%subtile_cols);
const uint32_t addr_2 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+8)*subtile_cols + col%subtile_cols);
const int swizzle_1 = blit ^ ((addr_1 % swizzle_repeat) >> 7) << 4;
const int swizzle_2 = blit ^ ((addr_2 % swizzle_repeat) >> 7) << 4;
move<U>::lds(tmp[0].x, (addr_1+ 0)^swizzle_1);
move<U>::lds(tmp[0].y, (addr_1+ 4)^swizzle_1);
move<U>::lds(tmp[2].x, (addr_1+32)^swizzle_1);
move<U>::lds(tmp[2].y, (addr_1+36)^swizzle_1);
move<U>::lds(tmp[1].x, (addr_2+ 0)^swizzle_2);
move<U>::lds(tmp[1].y, (addr_2+ 4)^swizzle_2);
move<U>::lds(tmp[3].x, (addr_2+32)^swizzle_2);
move<U>::lds(tmp[3].y, (addr_2+36)^swizzle_2);
dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
if(blit) {
#pragma unroll
for(int k = 0; k < 4; k++) {
dst.tiles[i][j].data[k] = T2{dst.tiles[i][j].data[k].y, dst.tiles[i][j].data[k].x};
}
}
}
else {
// handle the column-major layout
int row = (local_warpid*warp_height + i)*dst.tile_size_row + 2*(warp_laneid % 4);
int col = j*dst.tile_size_col + (warp_laneid / 4);
U2 tmp[4];
move<U>::lds(tmp[0].x, src.idx(shared_addr, {row+0, col+0}));
move<U>::lds(tmp[0].y, src.idx(shared_addr, {row+1, col+0}));
move<U>::lds(tmp[1].x, src.idx(shared_addr, {row+0, col+8}));
move<U>::lds(tmp[1].y, src.idx(shared_addr, {row+1, col+8}));
move<U>::lds(tmp[2].x, src.idx(shared_addr, {row+8, col+0}));
move<U>::lds(tmp[2].y, src.idx(shared_addr, {row+9, col+0}));
move<U>::lds(tmp[3].x, src.idx(shared_addr, {row+8, col+8}));
move<U>::lds(tmp[3].y, src.idx(shared_addr, {row+9, col+8}));
dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
}
}
}
}
/**
* @brief Collaboratively store data into a shared tile from register tiles split across a warpgroup.
*
* @tparam RT The register tile type
* @tparam ST The shared tile type
* @param dst[out] The destination shared tile.
* @param src[in] The source register tile.
*/
template<ducks::st::all ST, ducks::rt::all RT>
__device__ inline static void store(ST &dst, const RT &src) {
constexpr int height = ST::height;
constexpr int warp_height = RT::height;
static_assert(height%GROUP_WARPS == 0, "Group load / store requires tile height to be a multiple of GROUP_WARPS.");
static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
int local_warpid;
if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
else local_warpid = warpid();
using T2 = RT::dtype;
using U = ST::dtype;
using T = base_types::packing<T2>::unpacked_type;
using U2 = base_types::packing<U>::packed_type;
int warp_laneid = ::kittens::laneid();
// convert to shared state space
uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
#pragma unroll
for(int i = 0; i < warp_height; i++) {
#pragma unroll
for(int j = 0; j < src.width; j++) {
if constexpr (sizeof(typename ST::dtype) == 2) {
// handle the row-major layout
U2 tmp[4];
tmp[0] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
tmp[1] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
tmp[2] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
tmp[3] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
#ifdef KITTENS_HOPPER
int row = (local_warpid*warp_height + i)*src.tile_size_row + (warp_laneid % 16);
int col = j*src.tile_size_col + (warp_laneid / 16) * 8;
if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
move<U2>::stsm4(dst.idx(shared_addr, {row, col}), tmp[0], tmp[1], tmp[2], tmp[3]);
}
else {
move<U2>::stsm4t(dst.idx(shared_addr, {row, col}), tmp[0], tmp[2], tmp[1], tmp[3]);
}
#else
if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
int row = (local_warpid*warp_height + i)*src.tile_size_row + (warp_laneid / 4);
int col = j*src.tile_size_col + 2*(warp_laneid % 4);
move<U2>::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0]);
move<U2>::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[1]);
move<U2>::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[2]);
move<U2>::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3]);
}
else {
int row = (local_warpid*warp_height + i)*src.tile_size_row + 2*(warp_laneid % 4);
int col = j*src.tile_size_col + (warp_laneid / 4);
move<U>::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0].x);
move<U>::sts(dst.idx(shared_addr, {row+1, col+0}), tmp[0].y);
move<U>::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[1].x);
move<U>::sts(dst.idx(shared_addr, {row+1, col+8}), tmp[1].y);
move<U>::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[2].x);
move<U>::sts(dst.idx(shared_addr, {row+9, col+0}), tmp[2].y);
move<U>::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3].x);
move<U>::sts(dst.idx(shared_addr, {row+9, col+8}), tmp[3].y);
}
#endif
}
else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 1) {
// handle the row-major layout for 8-bit types
int warp_group_16 = (warp_laneid / 16); // divide each warp into two groups of 16 threads
int lane_in_16 = warp_laneid % 16; // position in group of 16 threads
int row = (local_warpid*warp_height + i)*src.tile_size_row + (lane_in_16 % 16); // find base row for warp in warpgroup and then distribute the 16 threads in the warp across the rows
int col = j*src.tile_size_col + warp_group_16 * 16; // find base column and then *16 for second half of the warp
U2 tmp[4];
tmp[0] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
tmp[1] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
tmp[2] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
tmp[3] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
move<U2>::stsm4(dst.idx(shared_addr, {row, col}), tmp[0], tmp[1], tmp[2], tmp[3]);
}
else {
move<U2>::stsm4t(dst.idx(shared_addr, {row, col}), tmp[0], tmp[2], tmp[1], tmp[3]);
}
}
else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 4) {
// handle the row-major layout for 32-bit types
int row = (local_warpid*warp_height + i)*src.tile_size_row + (warp_laneid / 4);
int col = j*src.tile_size_col + 2*(warp_laneid % 4);
if constexpr (ST::rows != ST::underlying_rows || ST::cols != ST::underlying_cols) { // subtile case
row += dst.row_offset;
col += dst.col_offset;
}
int blit = sizeof(typename ST::dtype) * ((warp_laneid%4) / 2);
T2 reg_tmp[4];
if(blit) {
#pragma unroll
for(int k = 0; k < 4; k++) {
reg_tmp[k] = T2{src.tiles[i][j].data[k].y, src.tiles[i][j].data[k].x};
}
}
else {
#pragma unroll
for(int k = 0; k < 4; k++) {
reg_tmp[k] = src.tiles[i][j].data[k];
}
}
U2 tmp[4];
tmp[0] = base_types::convertor<U2, T2>::convert(reg_tmp[0]);
tmp[1] = base_types::convertor<U2, T2>::convert(reg_tmp[1]);
tmp[2] = base_types::convertor<U2, T2>::convert(reg_tmp[2]);
tmp[3] = base_types::convertor<U2, T2>::convert(reg_tmp[3]);
static constexpr int swizzle_repeat = ST::swizzle_bytes * 8;
static constexpr int subtile_cols = ST::swizzle_bytes / sizeof(U);
const int outer_idx = col/subtile_cols;
const uint32_t addr_1 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+0)*subtile_cols + col%subtile_cols);
const uint32_t addr_2 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+8)*subtile_cols + col%subtile_cols);
const int swizzle_1 = blit ^ ((addr_1 % swizzle_repeat) >> 7) << 4;
const int swizzle_2 = blit ^ ((addr_2 % swizzle_repeat) >> 7) << 4;
move<U>::sts((addr_1+ 0)^swizzle_1, tmp[0].x);
move<U>::sts((addr_1+ 4)^swizzle_1, tmp[0].y);
move<U>::sts((addr_1+32)^swizzle_1, tmp[2].x);
move<U>::sts((addr_1+36)^swizzle_1, tmp[2].y);
move<U>::sts((addr_2+ 0)^swizzle_2, tmp[1].x);
move<U>::sts((addr_2+ 4)^swizzle_2, tmp[1].y);
move<U>::sts((addr_2+32)^swizzle_2, tmp[3].x);
move<U>::sts((addr_2+36)^swizzle_2, tmp[3].y);
}
else {
// handle the column-major layout
int row = (local_warpid*warp_height + i)*src.tile_size_row + 2*(warp_laneid % 4);
int col = j*src.tile_size_col + (warp_laneid / 4);
U2 tmp[4];
tmp[0] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
tmp[1] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
tmp[2] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
tmp[3] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
move<U>::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0].x);
move<U>::sts(dst.idx(shared_addr, {row+1, col+0}), tmp[0].y);
move<U>::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[1].x);
move<U>::sts(dst.idx(shared_addr, {row+1, col+8}), tmp[1].y);
move<U>::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[2].x);
move<U>::sts(dst.idx(shared_addr, {row+9, col+0}), tmp[2].y);
move<U>::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3].x);
move<U>::sts(dst.idx(shared_addr, {row+9, col+8}), tmp[3].y);
}
}
}
}
// Load and store of vectors from/to shared tiles.
template<ducks::rv::naive_layout RV, ducks::st::all ST>
__device__ inline static auto load(RV &dst, const ST &src, int2 row_col) {
KITTENS_CHECK_WARP;
static_assert(ST::cols>=RV::length, "Shared tile must be at least as wide as the vector.");
using T = RV::T;
using U = ST::T;
int warp_laneid = ::kittens::laneid();
// convert to shared state space
uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
#pragma unroll
for(int col = warp_laneid; col < dst.length; col+=WARP_THREADS) {
U tmp;
move<U>::lds(tmp, src.idx(shared_addr, {row_col.x, row_col.y + col}));
dst.data[col/WARP_THREADS][0] = base_types::convertor<T, U>::convert(tmp);
}
}
template<ducks::rv::naive_layout RV, ducks::st::all ST>
__device__ inline static auto store(ST &dst, const RV &src, int2 row_col) {
KITTENS_CHECK_WARP;
static_assert(ST::cols>=RV::length, "Shared tile must be at least as wide as the vector.");
using T = RV::T;
using U = ST::T;
int warp_laneid = ::kittens::laneid();
// convert to shared state space
uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
#pragma unroll
for(int col = warp_laneid; col < src.length; col+=WARP_THREADS) {
U tmp = base_types::convertor<U, T>::convert(src.data[col/WARP_THREADS][0]);
move<U>::sts(dst.idx(shared_addr, {row_col.x, row_col.y + col}), tmp);
}
}
@@ -1,325 +0,0 @@
/**
* @file
* @brief Group (collaborative warp) ops for loading tensor tiles into register tiles.
*/
/**
* @brief Load data from a tensor tile into a register tile.
*
* @tparam RT The register tile type
* @tparam TM The tensor memory tile type
* @param dst[out] The destination register tile.
* @param src[in] The source tensor tile.
*/
template<ducks::rt::row_layout RT, ducks::tt::all TM>
__device__ inline static void load_async(RT &dst, const TM &src) {
if constexpr (GROUP_WARPS == 1) {
static_assert(RT::height == TM::height, "register tile and tensor tile must match height");
static_assert(RT::width == TM::width, "register tile and tensor tile must match width");
using T2 = RT::dtype;
using U = typename TM::dtype;
using U2 = base_types::packing<typename TM::dtype>::packed_type;
if constexpr (sizeof(typename TM::dtype) == 1) {
#pragma unroll
for(int i = 0; i < dst.height; i++) {
#pragma unroll
for(int j = 0; j < dst.width; j++) {
asm volatile(
"tcgen05.ld.sync.aligned.16x128b.x2.pack::16b.b32 {%0, %1, %2, %3}, [%4];\n"
: "=r"(*(uint32_t*) &dst.tiles[i][j].data[0]),
"=r"(*(uint32_t*) &dst.tiles[i][j].data[1]),
"=r"(*(uint32_t*) &dst.tiles[i][j].data[2]),
"=r"(*(uint32_t*) &dst.tiles[i][j].data[3])
: "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
);
}
}
} else if constexpr (sizeof(typename TM::dtype) == 2) {
#pragma unroll
for(int i = 0; i < dst.height; i++) {
#pragma unroll
for(int j = 0; j < dst.width; j++) {
asm volatile(
"tcgen05.ld.sync.aligned.16x128b.x2.pack::16b.b32 {%0, %1, %2, %3}, [%4];\n"
: "=r"(*(uint32_t*) &dst.tiles[i][j].data[0]),
"=r"(*(uint32_t*) &dst.tiles[i][j].data[1]),
"=r"(*(uint32_t*) &dst.tiles[i][j].data[2]),
"=r"(*(uint32_t*) &dst.tiles[i][j].data[3])
: "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col))
);
}
}
}
else if constexpr (sizeof(typename TM::dtype) == 4) {
#pragma unroll
for(int i = 0; i < dst.height; i++) {
if constexpr (dst.width%4 == 0) {
#pragma unroll
for(int j = 0; j < dst.width; j+=4) {
U2 data[16];
asm volatile(
"tcgen05.ld.sync.aligned.16x256b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];\n"
: "=f"(data[0].x), "=f"(data[0].y),
"=f"(data[1].x), "=f"(data[1].y),
"=f"(data[2].x), "=f"(data[2].y),
"=f"(data[3].x), "=f"(data[3].y),
"=f"(data[4].x), "=f"(data[4].y),
"=f"(data[5].x), "=f"(data[5].y),
"=f"(data[6].x), "=f"(data[6].y),
"=f"(data[7].x), "=f"(data[7].y),
"=f"(data[8].x), "=f"(data[8].y),
"=f"(data[9].x), "=f"(data[9].y),
"=f"(data[10].x), "=f"(data[10].y),
"=f"(data[11].x), "=f"(data[11].y),
"=f"(data[12].x), "=f"(data[12].y),
"=f"(data[13].x), "=f"(data[13].y),
"=f"(data[14].x), "=f"(data[14].y),
"=f"(data[15].x), "=f"(data[15].y)
: "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
);
#pragma unroll
for(int k = 0; k < 4; k++) {
dst.tiles[i][j+0].data[k] = base_types::convertor<T2, U2>::convert(data[k]);
dst.tiles[i][j+1].data[k] = base_types::convertor<T2, U2>::convert(data[k+4]);
dst.tiles[i][j+2].data[k] = base_types::convertor<T2, U2>::convert(data[k+8]);
dst.tiles[i][j+3].data[k] = base_types::convertor<T2, U2>::convert(data[k+12]);
}
}
}
else if constexpr (dst.width%2 == 0) {
#pragma unroll
for(int j = 0; j < dst.width; j+=2) {
U2 data[8];
asm volatile(
"tcgen05.ld.sync.aligned.16x256b.x4.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];\n"
: "=f"(data[0].x), "=f"(data[0].y),
"=f"(data[1].x), "=f"(data[1].y),
"=f"(data[2].x), "=f"(data[2].y),
"=f"(data[3].x), "=f"(data[3].y),
"=f"(data[4].x), "=f"(data[4].y),
"=f"(data[5].x), "=f"(data[5].y),
"=f"(data[6].x), "=f"(data[6].y),
"=f"(data[7].x), "=f"(data[7].y)
: "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
);
#pragma unroll
for(int k = 0; k < 4; k++) {
dst.tiles[i][j+0].data[k] = base_types::convertor<T2, U2>::convert(data[k]);
dst.tiles[i][j+1].data[k] = base_types::convertor<T2, U2>::convert(data[k+4]);
}
}
}
else {
#pragma unroll
for(int j = 0; j < dst.width; j++) {
U2 data[4];
asm volatile(
"tcgen05.ld.sync.aligned.16x256b.x2.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n"
: "=f"(data[0].x), "=f"(data[0].y),
"=f"(data[1].x), "=f"(data[1].y),
"=f"(data[2].x), "=f"(data[2].y),
"=f"(data[3].x), "=f"(data[3].y)
: "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
);
#pragma unroll
for(int k = 0; k < 4; k++) {
dst.tiles[i][j].data[k] = base_types::convertor<T2, U2>::convert(data[k]);
}
}
}
}
}
}
else {
static_assert(GROUP_WARPS==4 || GROUP_WARPS==8);
constexpr int warp_rows = TM::rows/GROUP_WARPS;
static_assert(TM::cols==RT::cols);
static_assert(warp_rows==RT::rows);
if constexpr (GROUP_WARPS == 4) {
auto src_subtile = src.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*warpid(), 0);
::kittens::group<1>::load_async(dst, src_subtile);
}
else {
auto src_subtile = src.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*(warpid()%4)+16*(warpid()/4), 0);
::kittens::group<1>::load_async(dst, src_subtile);
}
}
}
/**
* @brief Store data into a tensor tile from a register tile.
*
* @tparam RT The register tile type
* @tparam TM The tensor memory tile type
* @param dst[out] The destination tensor tile.
* @param src[in] The source register tile.
*/
template<ducks::rt::all RT, ducks::tt::all TM>
__device__ inline static void store_async(TM &dst, const RT &src) {
if constexpr (GROUP_WARPS == 1) {
static_assert(RT::height == TM::height, "register tile and tensor tile must match height");
static_assert(RT::width == TM::width, "register tile and tensor tile must match width");
using T2 = RT::dtype;
using T = base_types::packing<T2>::unpacked_type;
using U = TM::dtype;
using U2 = base_types::packing<U>::packed_type;
if constexpr (sizeof(typename TM::dtype) == 2) {
#pragma unroll
for(int i = 0; i < src.height; i++) {
if constexpr (src.width%4 == 0) {
#pragma unroll
for(int j = 0; j < src.width; j+=4) {
asm volatile(
"tcgen05.st.sync.aligned.16x128b.x8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};\n"
:: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
"r"(*(uint32_t*)&src.tiles[i][j+0].data[0]),
"r"(*(uint32_t*)&src.tiles[i][j+0].data[1]),
"r"(*(uint32_t*)&src.tiles[i][j+0].data[2]),
"r"(*(uint32_t*)&src.tiles[i][j+0].data[3]),
"r"(*(uint32_t*)&src.tiles[i][j+1].data[0]),
"r"(*(uint32_t*)&src.tiles[i][j+1].data[1]),
"r"(*(uint32_t*)&src.tiles[i][j+1].data[2]),
"r"(*(uint32_t*)&src.tiles[i][j+1].data[3]),
"r"(*(uint32_t*)&src.tiles[i][j+2].data[0]),
"r"(*(uint32_t*)&src.tiles[i][j+2].data[1]),
"r"(*(uint32_t*)&src.tiles[i][j+2].data[2]),
"r"(*(uint32_t*)&src.tiles[i][j+2].data[3]),
"r"(*(uint32_t*)&src.tiles[i][j+3].data[0]),
"r"(*(uint32_t*)&src.tiles[i][j+3].data[1]),
"r"(*(uint32_t*)&src.tiles[i][j+3].data[2]),
"r"(*(uint32_t*)&src.tiles[i][j+3].data[3])
);
}
}
else if constexpr (src.width%2 == 0) {
#pragma unroll
for(int j = 0; j < src.width; j+=2) {
asm volatile(
"tcgen05.st.sync.aligned.16x128b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n"
:: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
"r"(*(uint32_t*)&src.tiles[i][j+0].data[0]),
"r"(*(uint32_t*)&src.tiles[i][j+0].data[1]),
"r"(*(uint32_t*)&src.tiles[i][j+0].data[2]),
"r"(*(uint32_t*)&src.tiles[i][j+0].data[3]),
"r"(*(uint32_t*)&src.tiles[i][j+1].data[0]),
"r"(*(uint32_t*)&src.tiles[i][j+1].data[1]),
"r"(*(uint32_t*)&src.tiles[i][j+1].data[2]),
"r"(*(uint32_t*)&src.tiles[i][j+1].data[3])
);
}
}
else {
#pragma unroll
for(int j = 0; j < src.width; j++) {
asm volatile(
"tcgen05.st.sync.aligned.16x128b.x2.b32 [%0], {%1, %2, %3, %4};\n"
:: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
"r"(*(uint32_t*)&src.tiles[i][j].data[0]),
"r"(*(uint32_t*)&src.tiles[i][j].data[1]),
"r"(*(uint32_t*)&src.tiles[i][j].data[2]),
"r"(*(uint32_t*)&src.tiles[i][j].data[3])
);
}
}
}
}
else if constexpr (sizeof(typename TM::dtype) == 4) {
#pragma unroll
for(int i = 0; i < src.height; i++) {
if constexpr(src.width%4 == 0) {
#pragma unroll
for(int j = 0; j < src.width; j+=4) {
U2 data[16];
#pragma unroll
for(int k = 0; k < 4; k++) {
data[k] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[k]);
data[k+4] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+1].data[k]);
data[k+8] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+2].data[k]);
data[k+12] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+3].data[k]);
}
asm volatile(
"tcgen05.st.sync.aligned.16x256b.x8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32};\n"
:: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
"f"(data[0].x), "f"(data[0].y),
"f"(data[1].x), "f"(data[1].y),
"f"(data[2].x), "f"(data[2].y),
"f"(data[3].x), "f"(data[3].y),
"f"(data[4].x), "f"(data[4].y),
"f"(data[5].x), "f"(data[5].y),
"f"(data[6].x), "f"(data[6].y),
"f"(data[7].x), "f"(data[7].y),
"f"(data[8].x), "f"(data[8].y),
"f"(data[9].x), "f"(data[9].y),
"f"(data[10].x), "f"(data[10].y),
"f"(data[11].x), "f"(data[11].y),
"f"(data[12].x), "f"(data[12].y),
"f"(data[13].x), "f"(data[13].y),
"f"(data[14].x), "f"(data[14].y),
"f"(data[15].x), "f"(data[15].y)
);
}
}
else if constexpr(src.width%2 == 0) {
#pragma unroll
for(int j = 0; j < src.width; j+=2) {
U2 data[8];
#pragma unroll
for(int k = 0; k < 4; k++) {
data[k] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[k]);
data[k+4] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+1].data[k]);
}
asm volatile(
"tcgen05.st.sync.aligned.16x256b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};\n"
:: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
"f"(data[0].x), "f"(data[0].y),
"f"(data[1].x), "f"(data[1].y),
"f"(data[2].x), "f"(data[2].y),
"f"(data[3].x), "f"(data[3].y),
"f"(data[4].x), "f"(data[4].y),
"f"(data[5].x), "f"(data[5].y),
"f"(data[6].x), "f"(data[6].y),
"f"(data[7].x), "f"(data[7].y)
);
}
}
else {
#pragma unroll
for(int j = 0; j < src.width; j++) {
U2 data[4];
#pragma unroll
for(int k = 0; k < 4; k++) {
data[k] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[k]);
}
asm volatile(
"tcgen05.st.sync.aligned.16x256b.x2.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n"
:: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
"f"(data[0].x), "f"(data[0].y),
"f"(data[1].x), "f"(data[1].y),
"f"(data[2].x), "f"(data[2].y),
"f"(data[3].x), "f"(data[3].y)
);
}
}
}
}
}
else {
static_assert(GROUP_WARPS==4 || GROUP_WARPS==8);
constexpr int warp_rows = TM::rows/GROUP_WARPS;
static_assert(TM::cols==RT::cols);
static_assert(warp_rows==RT::rows);
if constexpr (GROUP_WARPS == 4) {
auto dst_subtile = dst.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*warpid(), 0);
::kittens::group<1>::store_async(dst_subtile, src);
}
else {
auto dst_subtile = dst.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*(warpid()%4)+16*(warpid()/4), 0);
::kittens::group<1>::store_async(dst_subtile, src);
}
}
}
@@ -1,16 +0,0 @@
/**
* @file
* @brief An aggregate header of group memory operations on tiles.
*/
#include "shared_to_register.cuh"
#include "global_to_register.cuh"
#include "global_to_shared.cuh"
#ifdef KITTENS_BLACKWELL
#include "tensor_to_register.cuh"
#endif
#include "complex/complex_shared_to_register.cuh"
#include "complex/complex_global_to_register.cuh"
#include "complex/complex_global_to_shared.cuh"
@@ -1,134 +0,0 @@
/**
* @file
* @brief Functions for a group scope to call tile TMA functions.
*/
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::prefetch<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
}
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::prefetch<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx); // Don't do the mask
}
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
}
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
}
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
}
}
template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
}
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_add_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
}
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_add_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
}
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_add_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
}
}
template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_add_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
}
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_min_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
}
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_min_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
}
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_min_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
}
}
template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_min_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
}
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_max_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
}
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_max_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
}
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_max_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
}
}
template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) {
if(laneid() == 0) {
::kittens::tma::store_max_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
}
}
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) {
if(laneid() == 0) {
::kittens::tma::load_async<axis, policy, ST, GL, COORD>(dst, src, idx, bar); // Don't do the mask
}
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) {
if(laneid() == 0) {
::kittens::tma::load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar);
}
}
@@ -1,33 +0,0 @@
/**
* @file
* @brief Functions for a group scope to call tile TMA cluster functions.
*/
#ifdef KITTENS_BLACKWELL
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) {
if(laneid() == 0) {
::kittens::tma::cluster::load_async<axis, policy, ST, GL, COORD>(dst, src, idx, bar, cluster_mask, dst_mbar_cta);
}
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) {
if(laneid() == 0) {
::kittens::tma::cluster::load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar, cluster_mask, dst_mbar_cta);
}
}
#else
template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) {
if(laneid() == 0) {
::kittens::tma::cluster::load_async<axis, policy, ST, GL, COORD>(dst, src, idx, bar, cluster_mask);
}
}
template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) {
if(laneid() == 0) {
::kittens::tma::cluster::load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar, cluster_mask);
}
}
#endif
@@ -1,68 +0,0 @@
/**
* @file
* @brief Various utilities for group TMA memory operations.
*/
/* ---------- Barrier functions for async load ---------- */
/**
* @brief Sets the number of bytes expected at the semaphore.
*
* This function sets the number of bytes expected at the semaphore for the first thread in the warp.
* It converts the semaphore pointer to a generic shared memory pointer and uses an inline assembly
* instruction to set the expected number of bytes.
*
* @param semaphore Reference to the semaphore variable.
* @param bytes The number of bytes expected at the semaphore.
*/
__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes) {
if(laneid() == 0) {
::kittens::tma::expect_bytes(bar, bytes);
}
}
/**
* @brief Sets the number of bytes expected at the semaphore.
*
* This function sets the number of bytes expected at the mbarrier before the transaction arrives.
*/
template<typename T, typename... args>
__device__ static inline void expect(semaphore& bar, const T& _1, const args&... _2) {
expect_bytes(bar, size_bytes<T, args...>);
}
/* ---------- Synchronization functions for async store ---------- */
/**
* @brief Commits previous asynchronous TMA stores to a group and performs them.
*/
__device__ static inline void store_commit_group() {
asm volatile("cp.async.bulk.commit_group;");
}
/**
* @brief Waits for previous committed TMA store groups to complete.
*
* @tparam N The maximum number of remaining TMA store groups. Defaults to 0.
*/
template <int N=0>
__device__ static inline void store_async_wait() {
asm volatile (
"cp.async.bulk.wait_group %0;"
:
: "n"(N)
: "memory"
);
}
/**
* @brief Waits for previous committed TMA store groups to finish reading from shared memory.
*
* @tparam N The maximum number of remaining TMA store groups. Defaults to 0.
*/
template <int N=0>
__device__ static inline void store_async_read_wait() {
asm volatile (
"cp.async.bulk.wait_group.read %0;"
:
: "n"(N)
: "memory"
);
}
@@ -1,90 +0,0 @@
/**
* @brief Waits for the requested semaphore phase, at cluster scope
*
* @param semaphore Reference to the semaphore variable.
* @param kPhaseBit The phase bit used for the semaphore.
*/
__device__ static inline void wait(semaphore& bar, int kPhaseBit) {
void const* const ptr = &bar;
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
asm volatile (
"{\n"
".reg .pred P1;\n"
"LAB_WAIT:\n"
"mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n"
"@P1 bra.uni DONE;\n"
"bra.uni LAB_WAIT;\n"
"DONE:\n"
"}\n"
:: "r"(mbar_ptr),
"r"(kPhaseBit)
);
}
/**
* @brief Sets the number of bytes expected at the semaphore, assuming a multicast instruction.
*
* This function sets the number of bytes expected at the semaphore for the first thread in the warp.
* It converts the semaphore pointer to a generic shared memory pointer and uses an inline assembly
* instruction to set the expected number of bytes.
*
* It's worth being aware that this function is particularly necessary for multicast loads, and
* distributed shared memory can actually be done with a normal tma::expect followed by wait. See
* the unit tests of dsmem for an example.
*
* @param semaphore Reference to the semaphore variable.
* @param bytes The number of bytes expected at the semaphore.
*/
__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes, int dst_cta) {
if(laneid() == 0) {
::kittens::tma::cluster::expect_bytes(bar, bytes, dst_cta);
}
}
/**
* @brief Sets the number of bytes expected at the semaphore.
*
* This function sets the number of bytes expected at the semaphore for the first thread in the warp.
* It converts the semaphore pointer to a generic shared memory pointer and uses an inline assembly
* instruction to set the expected number of bytes.
*
* @tparam T The type of the data to be stored at the semaphore.
* @param semaphore Reference to the semaphore variable.
*/
/**
* @brief Sets the number of bytes expected at the semaphore.
*
* This function sets the number of bytes expected at the mbarrier before the transaction arrives.
*/
template<typename T, typename... args>
__device__ static inline void expect(semaphore& bar, int dst_cta, const T& _1, const args&... _2) {
expect_bytes(bar, size_bytes<T, args...>, dst_cta);
}
/**
* @brief Arrives at a semaphore in cluster scope.
*
* Marks a thread arrival at an mbarrier
*
* @param semaphore Reference to the semaphore variable.
* @param kPhaseBit The phase bit used for the semaphore.
*/
__device__ static inline void arrive(semaphore& bar, int dst_cta, uint32_t count=1) {
if(laneid() == 0) {
::kittens::tma::cluster::arrive(bar, dst_cta, count);
}
}
// Generic transfer
__device__ static inline void store_async(void *dst, void *src, int dst_cta, uint32_t size_bytes, semaphore& bar) {
if(laneid() == 0) {
::kittens::tma::cluster::store_async(dst, src, dst_cta, size_bytes, bar);
}
}
// Templated transfer for convenience
template<typename T>
__device__ static inline void store_async(T &dst_, T &src_, int dst_cta, semaphore& bar) {
store_async((void*)&dst_, (void*)&src_, dst_cta, size_bytes<T>, bar);
}
@@ -1,168 +0,0 @@
/**
* @file
* @brief Various utilities for group memory operations.
*/
template<int N=0> __device__ static inline void load_async_wait(int bar_id) { // for completing (non-TMA) async loads
asm volatile("cp.async.wait_group %0;\n" : : "n"(N) : "memory");
sync(bar_id);
}
template<int N=0> __device__ static inline void load_async_wait() { // for completing (non-TMA) async loads
KITTENS_CHECK_WARP
asm volatile("cp.async.wait_group %0;\n" : : "n"(N) : "memory");
__syncwarp();
}
__device__ static inline void arrive(barrier<GROUP_WARPS> bar) {
asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(GROUP_WARPS*WARP_THREADS) : "memory");
}
__device__ static inline void arrive_and_wait(barrier<GROUP_WARPS> bar) {
asm volatile("bar.sync %0, %1;\n" :: "r"(bar.barrier_id), "n"(GROUP_WARPS*WARP_THREADS) : "memory");
}
/**
* @brief Initializes a synchronization semaphore with a transaction count and sets the expected number of bytes.
*
* This function sets up a semaphore that is used to synchronize threads within a block during asynchronous operations.
* It initializes the semaphore with a thread count semaphore.
*
* Additionally, if it is given a shared tile type, it will also call `set_bytes` to prepare for the memory transaction.
*
* @param[out] semaphore The semaphore variable to initialize.
* @param[in] tc The thread counter for the semaphore.
*/
__device__ static inline void init_semaphore(semaphore& bar, int thread_count, int transaction_count=0) {
if (laneid() == 0) {
void const* const ptr = &bar;
uint32_t bar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
asm volatile (
"mbarrier.init.shared::cta.b64 [%0], %1;\n"
:: "r"(bar_ptr), "r"(thread_count+transaction_count)
);
}
}
/**
* @brief Invalidate an mbarrier
*
* @param[out] semaphore The semaphore variable to initialize.
* @param[in] tc The thread counter for the semaphore.
*/
__device__ static inline void invalidate_semaphore(semaphore& bar) {
if (laneid() == 0) {
void const* const ptr = &bar;
uint32_t bar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
asm volatile (
"mbarrier.inval.shared::cta.b64 [%0];\n"
:: "r"(bar_ptr)
);
}
}
/**
* @brief Arrives at a semaphore.
*
* Marks a warp arrival at an mbarrier
*
* @param semaphore Reference to the semaphore variable.
* @param kPhaseBit The phase bit used for the semaphore.
*/
__device__ static inline void arrive(semaphore& sem) {
if(laneid() == 0) {
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&sem));
asm volatile (
"mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n"
:
: "r"(mbar_ptr)
: "memory"
);
}
}
template<int num_warps> __device__ static inline void arrive(barrier<num_warps> bar) {
asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(num_warps*WARP_THREADS) : "memory");
}
#ifdef KITTENS_HOPPER
/**
* @brief Arrives at a semaphore.
*
* Marks a warp arrival at an mbarrier
*
* @param semaphore Reference to the semaphore variable.
* @param kPhaseBit The phase bit used for the semaphore.
*/
__device__ static inline void arrive(semaphore& sem, uint32_t count) {
if(laneid() == 0) {
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&sem));
asm volatile (
"mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n"
:
: "r"(mbar_ptr), "r"(count)
: "memory"
);
}
}
#endif
/**
* @brief Waits for the requested semaphore phase.
*
* @param semaphore Reference to the semaphore variable.
* @param kPhaseBit The phase bit used for the semaphore.
*/
__device__ static inline void wait(semaphore& sem, int kPhaseBit) {
void const* const ptr = &sem;
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
#ifdef KITTENS_HOPPER
asm volatile (
"{\n"
".reg .pred P1;\n"
"LAB_WAIT:\n"
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
"@P1 bra.uni DONE;\n"
"bra.uni LAB_WAIT;\n"
"DONE:\n"
"}\n"
:: "r"(mbar_ptr),
"r"(kPhaseBit)
);
#else
asm volatile (
"{\n"
".reg .pred P1;\n"
"LAB_WAIT:\n"
"mbarrier.test_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
"@P1 bra.uni DONE;\n"
"nanosleep.u32 5;\n" // wait a few nanoseconds on pre-Hopper architectures to save instruction issue slots
"bra.uni LAB_WAIT;\n"
"DONE:\n"
"}\n"
:: "r"(mbar_ptr),
"r"(kPhaseBit)
);
#endif
}
/**
* @brief Checks if the requested semaphore phase is ready.
*
* @param semaphore Reference to the semaphore variable.
* @param kPhaseBit The phase bit used for the semaphore.
*/
__device__ static inline int test_wait(semaphore& sem, int kPhaseBit) {
void const* const ptr = &sem;
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
int result;
asm volatile (
"{\n"
".reg .pred P1;\n"
"mbarrier.test_wait.parity.shared::cta.b64 P1, [%1], %2;\n"
"selp.u32 %0,1,0,P1;"
"}\n"
: "=r"(result)
: "r"(mbar_ptr), "r"(kPhaseBit)
);
return result;
}
@@ -1,138 +0,0 @@
/**
* @file
* @brief Functions for a warpgroup to collaboratively transfer data directly between global memory and registers and back.
*/
/**
* @brief Collaboratively loads data into register vectors from a source array in global memory.
*
* @tparam RV The register vector type.
* @tparam U The data type of the source array.
* @param[out] dst The destination register vector to load data into.
* @param[in] src The source array in global memory to load data from.
*/
template<ducks::rv::all RV, ducks::gl::all GL>
__device__ inline static void load(RV &dst, const GL &src, const coord<rv<typename RV::T, GROUP_WARPS*RV::length, typename RV::layout>> &idx) {
if constexpr (GROUP_WARPS == 1) {
using T2 = RV::dtype;
using U = typename GL::dtype;
using U2 = base_types::packing<U>::packed_type;
using T = base_types::packing<T2>::unpacked_type;
U *src_ptr = (U*)&src[(idx.template unit_coord<-1, 3>())];
int laneid = ::kittens::laneid();
if constexpr (std::is_same_v<typename RV::layout, align_l>) {
#pragma unroll
for(auto w = 0; w < (dst.outer_dim+3)/4; w++) {
int idx = w*64 + (laneid/4)*8 + 2*(laneid%4);
int o_dim = w*4 + (laneid/4) / 2;
int i_dim = (laneid/4) % 2;
// this should be a maximally coalesced load.
if(idx < dst.outer_dim*16)
dst[o_dim][i_dim] = base_types::convertor<T2, U2>::convert(*(U2*)&src_ptr[idx]);
}
// now we need to do a bunch of shuffle_sync's to make sure everyone has everything they need.
#pragma unroll
for(auto w = 0; w < dst.outer_dim; w++) {
int leader = 8*(w%4) + (laneid%4); // repeats every 64 columns
dst[w][0] = packed_shfl_sync(MASK_ALL, dst[w][0], leader);
dst[w][1] = packed_shfl_sync(MASK_ALL, dst[w][1], leader+4);
}
}
else if constexpr (std::is_same_v<typename RV::layout, ortho_l>) {
// really hoping https://stackoverflow.com/questions/15029765/is-coalescing-triggered-for-accessing-memory-in-reverse-order is still true
// otherwise there will be some pain :/
#pragma unroll
for(auto w = 0; w < (dst.outer_dim+1)/2; w++) {
int idx = w*32 + (laneid%4)*8 + (laneid/4);
int o_dim = w*2 + (laneid%4) / 2;
// this should be a maximally coalesced load.
if(idx < dst.outer_dim*16) {
T tmp = base_types::convertor<T, U>::convert(src_ptr[idx]);
if(laneid%2==0) dst[o_dim][0].x = tmp;
else dst[o_dim][0].y = tmp;
}
}
// now we need to do a bunch of shuffle_sync's to make sure everyone has everything they need.
#pragma unroll
for(auto w = 0; w < dst.outer_dim; w++) {
int leader = (laneid/4)*4 + 2*(w%2); // repeats every 64 columns
dst[w][0].x = __shfl_sync(MASK_ALL, dst[w][0].x, leader);
dst[w][0].y = __shfl_sync(MASK_ALL, dst[w][0].y, leader+1);
}
}
else if constexpr (std::is_same_v<typename RV::layout, naive_l>) {
#pragma unroll
for(auto w = 0; w < dst.outer_dim; w++) {
if(w < dst.outer_dim-1 || dst.length%32 == 0 || laneid<16) {
dst[w][0] = base_types::convertor<T, U>::convert(src_ptr[w*32 + laneid]);
}
}
}
}
else {
// Call warp level load
::kittens::group<1>::load(dst, src, coord<RV>(idx.b, idx.d, idx.r, idx.c*GROUP_WARPS+warpid()));
}
}
/**
* @brief Collaboratively stores data from register vectors to a destination array in global memory.
*
* @tparam RV The register vector type.
* @tparam U The data type of the destination array.
* @param[out] dst The destination array in global memory to store data into.
* @param[in] src The source register vector to store data from.
*/
template<ducks::rv::all RV, ducks::gl::all GL>
__device__ inline static void store(GL &dst, const RV &src, const coord<rv<typename RV::T, GROUP_WARPS*RV::length, typename RV::layout>> &idx) {
if constexpr (GROUP_WARPS == 1) {
using T2 = RV::dtype;
using U = typename GL::dtype;
using U2 = base_types::packing<U>::packed_type;
using T = base_types::packing<T2>::unpacked_type;
U *dst_ptr = (U*)&dst[(idx.template unit_coord<-1, 3>())];
int laneid = ::kittens::laneid();
if constexpr (std::is_same_v<typename RV::layout, align_l>) {
#pragma unroll
for(auto w = 0; w < (src.outer_dim+3)/4; w++) {
int idx = w*64 + (laneid/4)*8 + 2*(laneid%4);
int o_dim = w*4 + (laneid/4) / 2;
int i_dim = (laneid/4) % 2;
// this should be a maximally coalesced store. I hope!
if(idx < src.outer_dim*16)
*(U2*)&dst_ptr[idx] = base_types::convertor<U2, T2>::convert(src[o_dim][i_dim]);
}
}
else if constexpr (std::is_same_v<typename RV::layout, ortho_l>) {
// really hoping https://stackoverflow.com/questions/15029765/is-coalescing-triggered-for-accessing-memory-in-reverse-order is still true
// otherwise there will be some pain :/
#pragma unroll
for(auto w = 0; w < (src.outer_dim+1)/2; w++) {
int idx = w*32 + (laneid%4)*8 + (laneid/4);
int o_dim = w*2 + (laneid%4) / 2;
// this should be a maximally coalesced load.
if(idx < src.outer_dim*16) {
U tmp;
if(laneid%2==0) tmp = base_types::convertor<U, T>::convert(src[o_dim][0].x);
else tmp = base_types::convertor<U, T>::convert(src[o_dim][0].y);
dst_ptr[idx] = tmp;
}
}
}
else if constexpr (std::is_same_v<typename RV::layout, naive_l>) {
#pragma unroll
for(auto w = 0; w < src.outer_dim; w++) {
if(w < src.outer_dim-1 || src.length%32 == 0 || laneid<16) {
dst_ptr[w*32 + laneid] = base_types::convertor<U, T>::convert(src[w][0]);
}
}
}
}
else {
// Call warp level store
::kittens::group<1>::store(dst, src, coord<RV>(idx.b, idx.d, idx.r, idx.c*GROUP_WARPS+warpid()));
}
}
@@ -1,77 +0,0 @@
/**
* @file
* @brief Group (collaborative warp) ops for loading shared vectors from and storing to global memory.
*/
/**
* @brief Loads data from global memory into shared memory vector.
*
* This function loads data from a global memory location pointed to by `src` into a shared memory vector `dst`.
* It calculates the number of elements that can be transferred in one operation based on the size ratio of `float4` to the data type of `SV`.
* The function ensures coalesced memory access and efficient use of bandwidth by dividing the work among threads in a warp.
*
* @tparam SV Shared vector type, must satisfy ducks::sv::all concept.
* @param dst Reference to the shared vector where the data will be loaded.
* @param src Pointer to the global memory location from where the data will be loaded.
*/
template<ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void load(SV &dst, const GL &src, const COORD &idx) {
constexpr uint32_t elem_per_transfer = sizeof(float4) / sizeof(typename SV::dtype);
constexpr uint32_t total_calls = SV::length / elem_per_transfer; // guaranteed to divide
typename GL::dtype *src_ptr = (typename GL::dtype*)&src[(idx.template unit_coord<-1, 3>())];
uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
#pragma unroll
for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) {
if(i * elem_per_transfer < dst.length) {
float4 tmp;
move<float4>::ldg(tmp, (float4*)&src_ptr[i*elem_per_transfer]);
move<float4>::sts(dst_ptr + sizeof(typename SV::dtype)*i*elem_per_transfer, tmp);
}
}
}
/**
* @brief Stores data from a shared memory vector to global memory.
*
* This function stores data from a shared memory vector `src` to a global memory location pointed to by `dst`.
* Similar to the load function, it calculates the number of elements that can be transferred in one operation based on the size ratio of `float4` to the data type of `SV`.
* The function ensures coalesced memory access and efficient use of bandwidth by dividing the work among threads in a warp.
*
* @tparam SV Shared vector type, must satisfy ducks::sv::all concept.
* @param dst Pointer to the global memory location where the data will be stored.
* @param src Reference to the shared vector from where the data will be stored.
*/
template<ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store(GL &dst, const SV &src, const COORD &idx) {
constexpr uint32_t elem_per_transfer = sizeof(float4) / sizeof(typename SV::dtype);
constexpr uint32_t total_calls = SV::length / elem_per_transfer; // guaranteed to divide
typename GL::dtype *dst_ptr = (typename GL::dtype*)&dst[(idx.template unit_coord<-1, 3>())];
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
#pragma unroll
for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) {
if(i * elem_per_transfer < src.length) {
float4 tmp;
move<float4>::lds(tmp, src_ptr + sizeof(typename SV::dtype)*i*elem_per_transfer);
move<float4>::stg((float4*)&dst_ptr[i*elem_per_transfer], tmp);
}
}
}
template<ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx) {
constexpr uint32_t elem_per_transfer = sizeof(float4) / sizeof(typename SV::dtype);
constexpr uint32_t total_calls = SV::length / elem_per_transfer; // guaranteed to divide
typename GL::dtype *src_ptr = (typename GL::dtype*)&src[(idx.template unit_coord<-1, 3>())];
uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
#pragma unroll
for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) {
if(i * elem_per_transfer < dst.length) {
asm volatile(
"cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n"
:: "r"(dst_ptr + (uint32_t)sizeof(typename SV::dtype)*i*elem_per_transfer), "l"((uint64_t)&src_ptr[i*elem_per_transfer])
: "memory"
);
}
}
asm volatile("cp.async.commit_group;\n" ::: "memory");
}
@@ -1,159 +0,0 @@
/**
* @file
* @brief Functions for a group to collaboratively transfer data directly between shared memory and registers and back.
*/
/**
* @brief Collaboratively load data from a shared vector into register vectors split across a warpgroup.
*
* @tparam RV The register vector type
* @tparam SV The shared vector type
* @param dst[out] The destination register vector.
* @param src[in] The source shared vector.
*/
template<ducks::rv::all RV, ducks::sv::all SV>
__device__ inline static void load(RV &dst, const SV &src) {
using T2 = RV::dtype;
using U = SV::dtype;
using U2 = base_types::packing<U>::packed_type;
using T = base_types::packing<T2>::unpacked_type;
if constexpr (GROUP_WARPS == 1) {
static_assert(SV::length == RV::length);
int laneid = ::kittens::laneid();
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
__syncwarp();
if constexpr (std::is_same_v<typename RV::layout, align_l>) {
#pragma unroll
for(auto w = 0; w < (dst.outer_dim+3)/4; w++) {
int idx = w*64 + (laneid/4)*8 + 2*(laneid%4);
int o_dim = w*4 + (laneid/4) / 2;
int i_dim = (laneid/4) % 2;
// this should be a maximally coalesced load.
if(idx < dst.outer_dim*16) {
U2 tmp;
move<U2>::lds(tmp, src_ptr + sizeof(typename SV::dtype)*idx);
dst[o_dim][i_dim] = base_types::convertor<T2, U2>::convert(tmp);
}
}
__syncwarp();
// now we need to do a bunch of shuffle_sync's to make sure everyone has everything they need.
#pragma unroll
for(auto w = 0; w < dst.outer_dim; w++) {
int leader = 8*(w%4) + (laneid%4); // repeats every 64 columns
dst[w][0] = packed_shfl_sync(MASK_ALL, dst[w][0], leader);
dst[w][1] = packed_shfl_sync(MASK_ALL, dst[w][1], leader+4);
}
}
else if constexpr (std::is_same_v<typename RV::layout, ortho_l>) {
// really hoping https://stackoverflow.com/questions/15029765/is-coalescing-triggered-for-accessing-memory-in-reverse-order is still true
// otherwise there will be some pain :/
#pragma unroll
for(auto w = 0; w < (dst.outer_dim+1)/2; w++) {
int idx = w*32 + (laneid%4)*8 + (laneid/4);
int o_dim = w*2 + (laneid%4) / 2;
// this should be a maximally coalesced load.
if(idx < dst.outer_dim*16) {
U tmp;
move<U>::lds(tmp, src_ptr + sizeof(typename SV::dtype)*idx);
if(laneid%2==0) dst[o_dim][0].x = base_types::convertor<T, U>::convert(tmp);
else dst[o_dim][0].y = base_types::convertor<T, U>::convert(tmp);
}
}
__syncwarp();
// now we need to do a bunch of shuffle_sync's to make sure everyone has everything they need.
#pragma unroll
for(auto w = 0; w < dst.outer_dim; w++) {
int leader = (laneid/4)*4 + 2*(w%2); // repeats every 64 columns
dst[w][0].x = __shfl_sync(MASK_ALL, dst[w][0].x, leader);
dst[w][0].y = __shfl_sync(MASK_ALL, dst[w][0].y, leader+1);
}
}
else if constexpr (std::is_same_v<typename RV::layout, naive_l>) {
#pragma unroll
for(auto w = 0; w < dst.outer_dim; w++) {
if(w < dst.outer_dim-1 || RV::length%32 == 0 || laneid<16) {
U tmp;
move<U>::lds(tmp, src_ptr + sizeof(typename SV::dtype)*(w*32 + laneid));
dst[w][0] = base_types::convertor<T, U>::convert(tmp);
}
}
}
}
else {
static_assert(SV::length == RV::length*GROUP_WARPS);// confirm size correct
auto &_src = src.template subvec<RV::length>(warpid()); // pretend it's smaller and do warp-level load
::kittens::group<1>::load(dst, _src); // warp-level
}
}
/**
* @brief Collaboratively store data into a shared vector from register vectors split across a warpgroup.
*
* @tparam RV The register vector type
* @tparam SV The shared vector type
* @param dst[out] The destination shared vector.
* @param src[in] The source register vector.
*/
template<ducks::sv::all SV, ducks::rv::all RV>
__device__ inline static void store(SV &dst, const RV &src) {
using T2 = RV::dtype;
using U = SV::dtype;
using U2 = base_types::packing<U>::packed_type;
using T = base_types::packing<T2>::unpacked_type;
if constexpr (GROUP_WARPS == 1) {
static_assert(SV::length == RV::length);
int laneid = ::kittens::laneid();
uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
__syncwarp();
if constexpr (std::is_same_v<typename RV::layout, align_l>) {
#pragma unroll
for(auto w = 0; w < (src.outer_dim+3)/4; w++) {
int idx = w*64 + (laneid/4)*8 + 2*(laneid%4);
int o_dim = w*4 + (laneid/4) / 2;
int i_dim = (laneid/4) % 2;
// this should be a maximally coalesced store. I hope!
if(idx < src.outer_dim*16) {
U2 tmp = base_types::convertor<U2, T2>::convert(src[o_dim][i_dim]);
move<U2>::sts(dst_ptr + sizeof(typename SV::dtype)*idx, tmp);
}
}
}
else if constexpr (std::is_same_v<typename RV::layout, ortho_l>) {
// really hoping https://stackoverflow.com/questions/15029765/is-coalescing-triggered-for-accessing-memory-in-reverse-order is still true
// otherwise there will be some pain :/
#pragma unroll
for(auto w = 0; w < (src.outer_dim+1)/2; w++) {
int idx = w*32 + (laneid%4)*8 + (laneid/4);
int o_dim = w*2 + (laneid%4) / 2;
// this should be a maximally coalesced load.
if(idx < src.outer_dim*16) {
U tmp;
if(laneid%2==0) tmp = base_types::convertor<U, T>::convert(src[o_dim][0].x);
else tmp = base_types::convertor<U, T>::convert(src[o_dim][0].y);
move<U>::sts(dst_ptr + sizeof(typename SV::dtype)*idx, tmp);
}
}
}
else if constexpr (std::is_same_v<typename RV::layout, naive_l>) {
#pragma unroll
for(auto w = 0; w < src.outer_dim; w++) {
if(w < src.outer_dim-1 || RV::length%32 == 0 || laneid<16) {
U tmp = base_types::convertor<U, T>::convert(src[w][0]);
move<U>::sts(dst_ptr + sizeof(typename SV::dtype)*(w*32 + laneid), tmp);
}
}
}
}
else {
static_assert(SV::length == RV::length*GROUP_WARPS);// confirm size correct
auto &_dst = dst.template subvec<RV::length>(warpid()); // pretend it's smaller and do warp-level load
::kittens::group<1>::store(_dst, src); // warp-level
}
}
@@ -1,221 +0,0 @@
/**
* @file
* @brief Functions for a group scope to call vec TMA functions.
*/
/* ---------- Prefetch Tensor Map ---------- */
/**
* @brief Prefetches data from global memory into a shared memory vector, along with the tensormap.
*
* @tparam SV A shared vector type with a TMA-compatible layout
* @param[out] dst The destination shared memory vector.
* @param[in] src_tma_map The source tensormap address in global memory
* @param[in] vec_idx The coord of the requested vector.
*/
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void prefetch(SV &dst, const GL &src, const COORD &idx) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(src.template get_tma<SV, -1>());
for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
::kittens::detail::tma::vec_prefetch_tma_internal<policy>(tma_ptr, tma_coord);
}
}
__KITTENS_TMA_DEFINE_DEFAULT_LOAD_CACHE_VEC__(prefetch)
/* ---------- Async load and store data from gmem/smem ---------- */
/**
* @brief Asynchronously stores data into global memory from a shared memory vector.
*
* This function performs an asynchronous copy operation using CUDA's cp.async.bulk.tensor instruction.
*
* @tparam SV A shared vector type with a TMA-compatible layout
* @param[out] dst_tma_map The destination tensormap address in global memory
* @param[in] src The source shared memory vector.
* @param[in] vec_idx The coord of the vector destination.
*/
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_async(const GL &dst, const SV &src, const COORD &idx) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
store_commit_group();
}
__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_async)
template<cache_policy policy, ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_async(const PGL &dst, const SV &src, const COORD &idx) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
store_commit_group();
}
__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_async)
/**
* @brief Asynchronously performs an add reduction and stores the result into global memory.
*
* This function performs an asynchronous add reduction operation using CUDA's cp.reduce.async.bulk.tensor instruction.
*
* @tparam SV A shared vector type with a TMA-compatible layout
* @param[out] dst_tma_map The destination tensormap address in global memory
* @param[in] src The source shared memory vector.
* @param[in] vec_idx The coord of the vector destination.
*/
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_add_async(const GL &dst, const SV &src, const COORD &idx) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_add_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
store_commit_group();
}
__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_add_async)
template<cache_policy policy, ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_add_async(const PGL &dst, const SV &src, const COORD &idx) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_add_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
store_commit_group();
}
__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_add_async)
/**
* @brief Asynchronously performs an min reduction and stores the result into global memory.
*
* This function performs an asynchronous min reduction operation using CUDA's cp.reduce.async.bulk.tensor instruction.
*
* @tparam SV A shared vector type with a TMA-compatible layout
* @param[out] dst_tma_map The destination tensormap address in global memory
* @param[in] src The source shared memory vector.
* @param[in] vec_idx The coord of the vector destination.
*/
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_min_async(const GL &dst, const SV &src, const COORD &idx) {
static_assert(!std::is_same_v<typename SV::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_min_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
store_commit_group();
}
__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_min_async)
template<cache_policy policy, ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_min_async(const PGL &dst, const SV &src, const COORD &idx) {
static_assert(!std::is_same_v<typename SV::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_min_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
store_commit_group();
}
__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_min_async)
/**
* @brief Asynchronously performs an max reduction and stores the result into global memory.
*
* This function performs an asynchronous max reduction operation using CUDA's cp.reduce.async.bulk.tensor instruction.
*
* @tparam SV A shared vector type with a TMA-compatible layout
* @param[out] dst_tma_map The destination tensormap address in global memory
* @param[in] src The source shared memory vector.
* @param[in] vec_idx The coord of the vector destination.
*/
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_max_async(const GL &dst, const SV &src, const COORD &idx) {
static_assert(!std::is_same_v<typename SV::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_max_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
store_commit_group();
}
__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_max_async)
template<cache_policy policy, ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void store_max_async(const PGL &dst, const SV &src, const COORD &idx) {
static_assert(!std::is_same_v<typename SV::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_store_max_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
}
store_commit_group();
}
__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_max_async)
/**
* @brief Asynchronously loads data from global memory into a shared memory vector.
*
* This function performs an asynchronous copy operation using CUDA's cp.async.bulk.tensor instruction.
*
* @tparam SV A shared vector type with a TMA-compatible layout
* @param[out] dst The destination shared memory vector.
* @param[in] src_tma_map The source tensormap address in global memory
* @param[in] vec_idx The coord of the requested vector.
* @param[in,out] bar The semaphore used for synchronization of the asynchronous copy.
*/
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(src.template get_tma<SV, -1>());
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&bar));
uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst));
for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::vec_load_async_tma_internal<policy>(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord);
}
}
__KITTENS_TMA_DEFINE_SEMAPHORE_CACHE_VEC__(load_async)
@@ -1,31 +0,0 @@
/**
* @file
* @brief Functions for a group scope to call vec TMA cluster functions.
*/
/**
* @brief Asynchronously loads data from global memory into a shared memory vector, broadcast across a cluster
*
* This function performs an asynchronous copy operation using CUDA's cp.async.bulk.tensor instruction.
*
* @tparam SV A shared vector type with a TMA-compatible layout
* @param[out] dst The destination shared memory vector.
* @param[in] src_tma_map The source tensormap address in global memory
* @param[in,out] bar The semaphore used for synchronization of the asynchronous copy.
* @param[in] vec_idx The coord of the requested vector.
* @param[in] cluster_mask The mask of the clusters to broadcast to.
*/
template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) {
coord<> unit_coord = idx.template unit_coord<-1, 3>();
uint64_t tma_ptr = reinterpret_cast<uint64_t>(src.template get_tma<SV, -1>());
uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&bar));
uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst));
for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
coord<> tma_coord = unit_coord;
tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
::kittens::detail::tma::cluster::vec_load_async_tma_internal<policy>(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord, cluster_mask, dst_mbar_cta);
}
}
__KITTENS_TMA_DEFINE_CLUSTER_SEMAPHORE_CACHE_VEC__(load_async)
@@ -1,8 +0,0 @@
/**
* @file
* @brief An aggregate header of group memory operations on vectors.
*/
#include "shared_to_register.cuh"
#include "global_to_register.cuh"
#include "global_to_shared.cuh"
@@ -1,17 +0,0 @@
/**
* @file
* @brief An aggregate header for all group-scope MMA operations.
*/
// All compilation targets can use the warp-scope MMA operations.
#include "warp/warp.cuh"
// Hopper has its own warpgroup-scope MMA operations.
#ifdef KITTENS_HOPPER
#include "warpgroup/warpgroup.cuh"
#endif
// Blackwell has its own tensor-scope MMA operations.
#ifdef KITTENS_BLACKWELL
#include "tensor/tensor.cuh"
#endif
@@ -1,172 +0,0 @@
/**
* @file Group-level tcgen05 MMA operations.
*/
template<int trans_a, int n_trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B, int acc=1, int ncta=1>
__device__ static inline void mma(D &d, const A &a, const B &b, semaphore &sem) {
if(laneid() == 0) ::kittens::mma<trans_a, n_trans_b, D, A, B, acc, ncta>(d, a, b, sem);
}
template<int trans_a, int trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B, int acc=1>
__device__ static inline void mma2(D &d, const A &a, const B &b, semaphore &sem) {
mma<trans_a, trans_b, D, A, B, acc, 2>(d, a, b, sem);
}
template<int trans_a, int trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm(D &d, const A &a, const B &b, semaphore &sem) {
mma<trans_a, trans_b, D, A, B, 0>(d, a, b, sem);
}
template<int trans_a, int trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2(D &d, const A &a, const B &b, semaphore &sem) {
mma2<trans_a, trans_b, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_AB(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::N, transpose::N, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_AB(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::N, transpose::N, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_ABt(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::N, transpose::T, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_ABt(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::N, transpose::T, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_AtB(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::T, transpose::N, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_AtB(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::T, transpose::N, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_AtBt(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::T, transpose::T, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_AtBt(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::T, transpose::T, D, A, B, 1>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_AB(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::N, transpose::N, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_AB(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::N, transpose::N, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_ABt(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::N, transpose::T, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_ABt(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::N, transpose::T, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_AtB(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::T, transpose::N, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_AtB(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::T, transpose::N, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_AtBt(D &d, const A &a, const B &b, semaphore &sem) {
mma<transpose::T, transpose::T, D, A, B, 0>(d, a, b, sem);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_AtBt(D &d, const A &a, const B &b, semaphore &sem) {
mma2<transpose::T, transpose::T, D, A, B, 0>(d, a, b, sem);
}
// no sem versions
template<int trans_a, int n_trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B, int acc=1, int ncta=1>
__device__ static inline void mma(D &d, const A &a, const B &b) {
if(laneid() == 0) ::kittens::mma<trans_a, n_trans_b, D, A, B, acc, ncta>(d, a, b);
}
template<int trans_a, int trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B, int acc=1>
__device__ static inline void mma2(D &d, const A &a, const B &b) {
mma<trans_a, trans_b, D, A, B, acc, 2>(d, a, b);
}
template<int trans_a, int trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm(D &d, const A &a, const B &b) {
mma<trans_a, trans_b, D, A, B, 0>(d, a, b);
}
template<int trans_a, int trans_b, ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2(D &d, const A &a, const B &b) {
mma2<trans_a, trans_b, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_AB(D &d, const A &a, const B &b) {
mma<transpose::N, transpose::N, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_AB(D &d, const A &a, const B &b) {
mma2<transpose::N, transpose::N, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_ABt(D &d, const A &a, const B &b) {
mma<transpose::N, transpose::T, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_ABt(D &d, const A &a, const B &b) {
mma2<transpose::N, transpose::T, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_AtB(D &d, const A &a, const B &b) {
mma<transpose::T, transpose::N, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_AtB(D &d, const A &a, const B &b) {
mma2<transpose::T, transpose::N, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma_AtBt(D &d, const A &a, const B &b) {
mma<transpose::T, transpose::T, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mma2_AtBt(D &d, const A &a, const B &b) {
mma2<transpose::T, transpose::T, D, A, B, 1>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_AB(D &d, const A &a, const B &b) {
mma<transpose::N, transpose::N, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_AB(D &d, const A &a, const B &b) {
mma2<transpose::N, transpose::N, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_ABt(D &d, const A &a, const B &b) {
mma<transpose::N, transpose::T, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_ABt(D &d, const A &a, const B &b) {
mma2<transpose::N, transpose::T, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_AtB(D &d, const A &a, const B &b) {
mma<transpose::T, transpose::N, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_AtB(D &d, const A &a, const B &b) {
mma2<transpose::T, transpose::N, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm_AtBt(D &d, const A &a, const B &b) {
mma<transpose::T, transpose::T, D, A, B, 0>(d, a, b);
}
template<ducks::tt::all D, typename A, ducks::st_descriptor::input B>
__device__ static inline void mm2_AtBt(D &d, const A &a, const B &b) {
mma2<transpose::T, transpose::T, D, A, B, 0>(d, a, b);
}
@@ -1,947 +0,0 @@
/**
* @file
* @brief Matrix multiply-accumulate operations for tiles stored in registers.
*/
/**
* @brief Perform the HMMA.16816 operation.
*
* This function performs the half-precision matrix multiply-accumulate operation
* using the `mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32` instruction.
*
* @param[out] d0 The first half of the output float2 accumulator.
* @param[out] d1 The second half of the output float2 accumulator.
* @param[in] a0 The first half of the first input bf16_2 matrix.
* @param[in] a1 The second half of the first input bf16_2 matrix.
* @param[in] a2 The first half of the second input bf16_2 matrix.
* @param[in] a3 The second half of the second input bf16_2 matrix.
* @param[in] b0 The first half of the bf16_2 matrix B.
* @param[in] b1 The second half of the bf16_2 matrix B.
* @param[in] c0 The first half of the float2 accumulator matrix C.
* @param[in] c1 The second half of the float2 accumulator matrix C.
*/
__device__ static inline void hmma16816( float2 &d0, float2 &d1,
const bf16_2 &a0, const bf16_2 &a1, const bf16_2 &a2, const bf16_2 &a3,
const bf16_2 &b0, const bf16_2 &b1,
const float2 &c0, const float2 &c1 ) {
asm volatile(
// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#multiply-and-accumulate-instruction-mma
"mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " \
"{%0, %1, %2, %3}, " \
"{%4, %5, %6, %7}, " \
"{%8, %9}, " \
"{%10, %11, %12, %13};"
// D matrix
: "+f"(d0.x), "+f"(d0.y),
"+f"(d1.x), "+f"(d1.y)
// A matrix
: "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)),
"r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)),
// B matrix
"r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)),
// C matrix
"f"(c0.x), "f"(c0.y),
"f"(c1.x), "f"(c1.y)
);
}
/**
* @brief Perform the HMMA.16816 operation with inputs as fp16 and fp32 accumulators
*
* This function performs the half-precision matrix multiply-accumulate operation
* using the `mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32` instruction.
*
* @param[out] d0 The first half of the output float2 accumulator.
* @param[out] d1 The second half of the output float2 accumulator.
* @param[in] a0 The first half of the first input half_2 matrix.
* @param[in] a1 The second half of the first input half_2 matrix.
* @param[in] a2 The first half of the second input half_2 matrix.
* @param[in] a3 The second half of the second input half_2 matrix.
* @param[in] b0 The first half of the half_2 matrix B.
* @param[in] b1 The second half of the half_2 matrix B.
* @param[in] c0 The first half of the float2 accumulator matrix C.
* @param[in] c1 The second half of the float2 accumulator matrix C.
*/
__device__ static inline void hmma16816( float2 &d0, float2 &d1,
const half_2 &a0, const half_2 &a1, const half_2 &a2, const half_2 &a3,
const half_2 &b0, const half_2 &b1,
const float2 &c0, const float2 &c1 ) {
asm volatile(
// https://docs.nvidia.com/cuda/parallel-thread-execution/#multiply-and-accumulate-instruction-mma
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " \
"{%0, %1, %2, %3}, " \
"{%4, %5, %6, %7}, " \
"{%8, %9}, " \
"{%10, %11, %12, %13};"
// D matrix
: "+f"(d0.x), "+f"(d0.y),
"+f"(d1.x), "+f"(d1.y)
// A matrix
: "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)),
"r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)),
// B matrix
"r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)),
// C matrix
"f"(c0.x), "f"(c0.y),
"f"(c1.x), "f"(c1.y)
);
}
/**
* @brief Perform the HMMA.16816 operation.
*
* This function performs the half-precision matrix multiply-accumulate operation
* using the `mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16` instruction.
*
* @param[out] d0 The first half of the output half_2 accumulator.
* @param[out] d1 The second half of the output half_2 accumulator.
* @param[in] a0 The first half of the first input half_2 matrix.
* @param[in] a1 The second half of the first input half_2 matrix.
* @param[in] a2 The first half of the second input half_2 matrix.
* @param[in] a3 The second half of the second input half_2 matrix.
* @param[in] b0 The first half of the half_2 matrix B.
* @param[in] b1 The second half of the half_2 matrix B.
* @param[in] c0 The first half of the half_2 accumulator matrix C.
* @param[in] c1 The second half of the half_2 accumulator matrix C.
*/
__device__ static inline void hmma16816( half_2 &d0, half_2 &d1,
const half_2 &a0, const half_2 &a1, const half_2 &a2, const half_2 &a3,
const half_2 &b0, const half_2 &b1,
const half_2 &c0, const half_2 &c1 ) {
asm volatile(
// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#multiply-and-accumulate-instruction-mma
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " \
"{%0, %1}, " \
"{%2, %3, %4, %5}, " \
"{%6, %7}, " \
"{%8, %9};"
// D matrix
: "=r"(*(uint32_t*)(&d0)), "=r"(*(uint32_t*)(&d1))
// A matrix
: "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)),
"r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)),
// B matrix
"r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)),
// C matrix
"r"(*(uint32_t*)(&c0)), "r"(*(uint32_t*)(&c1))
);
}
#ifdef KITTENS_HOPPER
/**
* @brief Perform the HMMA.16816 operation for FP8 using fp8e4m3_2.
*
* Using mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 instruction
* but with fp8e4m3_2 (2 FP8 values) instead of fp8e4m3_4
*/
/**
* @brief Perform the HMMA.16816 operation for FP8.
*
* This function performs the fp8-precision matrix multiply-accumulate operation
* using the `mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32` instruction.
*
* @param[out] d0 The first half of the output float2 accumulator.
* @param[out] d1 The second half of the output float2 accumulator.
* @param[in] a0,a1,a2,a3 Input FP8 matrix A values
* @param[in] b0,b1 Input FP8 matrix B values
* @param[in] c0,c1 Input float2 accumulator matrix C values
*/
__device__ static inline void hmma16816( float2 &d0, float2 &d1,
const fp8e4m3_4 &a0, const fp8e4m3_4 &a1,
const fp8e4m3_4 &a2, const fp8e4m3_4 &a3,
const fp8e4m3_4 &b0, const fp8e4m3_4 &b1,
const float2 &c0, const float2 &c1) {
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
"{%0, %1, %2, %3}, "
"{%4, %5, %6, %7}, "
"{%8, %9}, "
"{%10, %11, %12, %13};"
// D matrix (output)
: "+f"(d0.x), "+f"(d0.y),
"+f"(d1.x), "+f"(d1.y)
// A matrix
: "r"(*(uint32_t*)(&a0)), "r"(*(uint32_t*)(&a1)),
"r"(*(uint32_t*)(&a2)), "r"(*(uint32_t*)(&a3)),
// B matrix
"r"(*(uint32_t*)(&b0)), "r"(*(uint32_t*)(&b1)),
// C matrix
"f"(c0.x), "f"(c0.y),
"f"(c1.x), "f"(c1.y)
);
}
#endif
/**
* @brief Base matrix multiply-accumulate operation for row layout.
*
* This function performs the base matrix multiply-accumulate operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<float2, row_layout> accumulator.
* @param[in] a The first input rt_base<bf16_2, row_layout> matrix.
* @param[in] b The second input rt_base<bf16_2, col_layout> matrix in column-major mode.
* @param[in] c The input rt_base<float2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_AB_base(rt_base<float, ducks::rt_layout::row> &d,
const rt_base<bf16, ducks::rt_layout::row> &a,
const rt_base<bf16, ducks::rt_layout::col> &b, // in col-major mode
const rt_base<float, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2],
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3],
c.data[2], c.data[3]
);
}
/**
* @brief Base matrix multiply-accumulate operation for row layout
* with fp16 inputs and fp32 accumulators.
*
* This function performs the base matrix multiply-accumulate operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<float2, row_layout> accumulator.
* @param[in] a The first input rt_base<half_2, row_layout> matrix.
* @param[in] b The second input rt_base<half_2, col_layout> matrix in column-major mode.
* @param[in] c The input rt_base<float2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_AB_base(rt_base<float, ducks::rt_layout::row> &d,
const rt_base<half, ducks::rt_layout::row> &a,
const rt_base<half, ducks::rt_layout::col> &b, // in col-major mode
const rt_base<float, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2],
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3],
c.data[2], c.data[3]
);
}
#ifdef KITTENS_HOPPER
/**
* @brief Base matrix multiply-accumulate operation for row layout.
*
* This function performs the base matrix multiply-accumulate operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<float2, row_layout> accumulator.
* @param[in] a The first input rt_base<fp8e4m3, row_layout> matrix.
* @param[in] b The second input rt_base<fp8e4m3, col_layout> matrix in column-major mode.
* @param[in] c The input rt_base<float2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_AB_base(rt_base<float, ducks::rt_layout::row> &d,
const rt_base<fp8e4m3, ducks::rt_layout::row> &a,
const rt_base<fp8e4m3, ducks::rt_layout::col> &b, // in col-major mode
const rt_base<float, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2],
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3],
c.data[2], c.data[3]
);
}
#endif
/**
* @brief Base matrix multiply-accumulate operation for row layout.
*
* This function performs the base matrix multiply-accumulate operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<half_2, row_layout> accumulator.
* @param[in] a The first input rt_base<half_2, row_layout> matrix.
* @param[in] b The second input rt_base<half_2, col_layout> matrix in column-major mode.
* @param[in] c The input rt_base<half_2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_AB_base(rt_base<half, ducks::rt_layout::row> &d,
const rt_base<half, ducks::rt_layout::row> &a,
const rt_base<half, ducks::rt_layout::col> &b, // in col-major mode
const rt_base<half, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2],
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3],
c.data[2], c.data[3]
);
}
/**
* @brief Base dot product operation for row layout.
*
* This function performs the base dot product operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<float2, row_layout> accumulator.
* @param[in] a The first input rt_base<bf16_2, row_layout> matrix.
* @param[in] b The second input rt_base<bf16_2, row_layout> matrix in row-major mode.
* @param[in] c The input rt_base<float2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_ABt_base(rt_base<float, ducks::rt_layout::row> &d,
const rt_base<bf16, ducks::rt_layout::row> &a,
const rt_base<bf16, ducks::rt_layout::row> &b, // in row-major mode
const rt_base<float, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2], // for some reason this one seems to need to be backwards
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3], // for some reason this one seems to need to be backwards
c.data[2], c.data[3]
);
}
/**
* @brief Base dot product operation for row layout
* with fp16 inputs and fp32 accumulators.
*
* This function performs the base dot product operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<float2, row_layout> accumulator.
* @param[in] a The first input rt_base<half_2, row_layout> matrix.
* @param[in] b The second input rt_base<half_2, row_layout> matrix in row-major mode.
* @param[in] c The input rt_base<float2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_ABt_base(rt_base<float, ducks::rt_layout::row> &d,
const rt_base<half, ducks::rt_layout::row> &a,
const rt_base<half, ducks::rt_layout::row> &b, // in row-major mode
const rt_base<float, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2], // for some reason this one seems to need to be backwards
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3], // for some reason this one seems to need to be backwards
c.data[2], c.data[3]
);
}
#ifdef KITTENS_HOPPER
/**
* @brief Base dot product operation for row layout.
*
* This function performs the base dot product operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<float2, row_layout> accumulator.
* @param[in] a The first input rt_base<fp8e4m3x4, row_layout> matrix.
* @param[in] b The second input rt_base<fp8e4m3x4, row_layout> matrix in row-major mode.
* @param[in] c The input rt_base<float2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_ABt_base(rt_base<float, ducks::rt_layout::row> &d,
const rt_base<fp8e4m3, ducks::rt_layout::row> &a,
const rt_base<fp8e4m3, ducks::rt_layout::row> &b, // in row-major mode
const rt_base<float, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2], // for some reason this one seems to need to be backwards
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3], // for some reason this one seems to need to be backwards
c.data[2], c.data[3]
);
}
#endif
/**
* @brief Base matrix multiply-accumulate operation for row layout with transposed A.
*
* This function performs the base matrix multiply-accumulate operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<float2, row_layout> accumulator.
* @param[in] a The first input rt_base<bf16_2, col_layout> matrix.
* @param[in] b The second input rt_base<bf16_2, col_layout> matrix in column-major mode.
* @param[in] c The input rt_base<float2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_AtB_base(rt_base<float, ducks::rt_layout::row> &d,
const rt_base<bf16, ducks::rt_layout::col> &a,
const rt_base<bf16, ducks::rt_layout::col> &b, // in col-major mode
const rt_base<float, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2],
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3],
c.data[2], c.data[3]
);
}
/**
* @brief Base matrix multiply-accumulate operation for row layout with transposed A
* with fp16 inputs and fp32 accumulators.
*
* This function performs the base matrix multiply-accumulate operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<float2, row_layout> accumulator.
* @param[in] a The first input rt_base<half_2, col_layout> matrix.
* @param[in] b The second input rt_base<half_2, col_layout> matrix in column-major mode.
* @param[in] c The input rt_base<float2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_AtB_base(rt_base<float, ducks::rt_layout::row> &d,
const rt_base<half, ducks::rt_layout::col> &a,
const rt_base<half, ducks::rt_layout::col> &b, // in col-major mode
const rt_base<float, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2],
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3],
c.data[2], c.data[3]
);
}
#ifdef KITTENS_HOPPER
/**
* @brief Base matrix multiply-accumulate operation for row layout with transposed A.
*
* This function performs the base matrix multiply-accumulate operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<float2, row_layout> accumulator.
* @param[in] a The first input rt_base<fp8e4m3x4, col_layout> matrix.
* @param[in] b The second input rt_base<fp8e4m3x4, col_layout> matrix in column-major mode.
* @param[in] c The input rt_base<float2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_AtB_base(rt_base<float, ducks::rt_layout::row> &d,
const rt_base<fp8e4m3, ducks::rt_layout::col> &a,
const rt_base<fp8e4m3, ducks::rt_layout::col> &b, // in col-major mode
const rt_base<float, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2],
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3],
c.data[2], c.data[3]
);
}
#endif
/**
* @brief Base matrix multiply-accumulate operation for row layout with transposed A and B.
*
* This function performs the base matrix multiply-accumulate operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<float2, row_layout> accumulator.
* @param[in] a The first input rt_base<bf16_2, col_layout> matrix.
* @param[in] b The second input rt_base<bf16_2, col_layout> matrix in column-major mode.
* @param[in] c The input rt_base<float2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_AtBt_base(rt_base<float, ducks::rt_layout::row> &d,
const rt_base<bf16, ducks::rt_layout::col> &a,
const rt_base<bf16, ducks::rt_layout::row> &b, // in col-major mode
const rt_base<float, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2],
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3],
c.data[2], c.data[3]
);
}
/**
* @brief Base matrix multiply-accumulate operation for row layout with transposed A and B
* with fp16 inputs and fp32 accumulators.
*
* This function performs the base matrix multiply-accumulate operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<float2, row_layout> accumulator.
* @param[in] a The first input rt_base<half_2, col_layout> matrix.
* @param[in] b The second input rt_base<half_2, row_layout> matrix in row-major mode.
* @param[in] c The input rt_base<float2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_AtBt_base(rt_base<float, ducks::rt_layout::row> &d,
const rt_base<half, ducks::rt_layout::col> &a,
const rt_base<half, ducks::rt_layout::row> &b, // in row-major mode
const rt_base<float, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2],
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3],
c.data[2], c.data[3]
);
}
#ifdef KITTENS_HOPPER
/**
* @brief Base matrix multiply-accumulate operation for row layout with transposed A and B.
*
* This function performs the base matrix multiply-accumulate operation
* using the `hmma16816` function for matrices in row layout.
*
* @param[out] d The output rt_base<float2, row_layout> accumulator.
* @param[in] a The first input rt_base<fp8e4m3x4, col_layout> matrix.
* @param[in] b The second input rt_base<fp8e4m3x4, col_layout> matrix in column-major mode.
* @param[in] c The input rt_base<float2, row_layout> accumulator matrix.
*/
__device__ static inline void mma_AtBt_base(rt_base<float, ducks::rt_layout::row> &d,
const rt_base<fp8e4m3, ducks::rt_layout::col> &a,
const rt_base<fp8e4m3, ducks::rt_layout::row> &b, // in col-major mode
const rt_base<float, ducks::rt_layout::row> &c) {
hmma16816(
d.data[0], d.data[1],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[0], b.data[2],
c.data[0], c.data[1]
);
hmma16816(
d.data[2], d.data[3],
a.data[0], a.data[1], a.data[2], a.data[3],
b.data[1], b.data[3],
c.data[2], c.data[3]
);
}
#endif
/**
* @brief Matrix multiply-accumulate operation.
*
* This function performs the matrix multiply-accumulate operation
* using the `hmma16816` function.
*
* @tparam N The number of row tiles.
* @tparam K The number of column tiles for the A matrix and row tiles for the B matrix.
* @tparam M The number of column tiles for the B matrix.
* @param[out] d The output rt_hf<N, M, row_layout> accumulator.
* @param[in] a The first input rt_hf<N, K, row_layout> matrix.
* @param[in] b The second input rt_hf<K, M, col_layout> matrix in column-major mode.
* @param[in] c The input rt_hf<N, M, row_layout> accumulator matrix.
*/
template<ducks::rt::row_layout D, ducks::rt::row_layout A, ducks::rt::col_layout B, ducks::rt::row_layout C>
__device__ static inline void mma_AB(D &d,
const A &a,
const B &b,
const C &c) {
KITTENS_CHECK_WARP
static_assert(D::rows == A::rows && D::cols == B::cols); // Check D matches A, B
static_assert(A::cols == B::rows); // Check reduction dim is same
static_assert(D::rows == C::rows && D::cols == C::cols); // Check D matches C
#ifdef KITTENS_HOPPER
static_assert(
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, bf16> &&
std::is_same_v<typename B::T, bf16> && std::is_same_v<typename C::T, float>) ||
(std::is_same_v<typename D::T, half> && std::is_same_v<typename A::T, half> &&
std::is_same_v<typename B::T, half> && std::is_same_v<typename C::T, half>) ||
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, fp8e4m3> &&
std::is_same_v<typename B::T, fp8e4m3> && std::is_same_v<typename C::T, float>)
);
#else
static_assert(
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, bf16> &&
std::is_same_v<typename B::T, bf16> && std::is_same_v<typename C::T, float>) ||
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, half> &&
std::is_same_v<typename B::T, half> && std::is_same_v<typename C::T, float>) ||
(std::is_same_v<typename D::T, half> && std::is_same_v<typename A::T, half> &&
std::is_same_v<typename B::T, half> && std::is_same_v<typename C::T, half>)
);
#endif
#pragma unroll
for(int n = 0; n < D::height; n++) {
#pragma unroll
for(int m = 0; m < D::width; m++) {
mma_AB_base(
d.tiles[n][m],
a.tiles[n][0],
b.tiles[0][m],
c.tiles[n][m]
);
#pragma unroll
for(int k = 1; k < A::width; k++) {
mma_AB_base(
d.tiles[n][m],
a.tiles[n][k],
b.tiles[k][m],
d.tiles[n][m]
);
}
}
}
}
/**
* @brief Dot product operation for row layout.
*
* This function performs the dot product operation
* using the `hmma16816` function.
*
* @tparam N The number of row tiles.
* @tparam K The number of column tiles for the A matrix and row tiles for the B matrix.
* @tparam M The number of column tiles for the B matrix.
* @param[out] d The output rt_fl<N, M, row_layout> accumulator.
* @param[in] a The first input rt_bf<N, K, row_layout> matrix.
* @param[in] b The second input rt_bf<M, K, row_layout> matrix in row-major mode.
* @param[in] c The input rt_fl<N, M, row_layout> accumulator matrix.
*/
template<ducks::rt::row_layout D, ducks::rt::row_layout A, ducks::rt::row_layout B, ducks::rt::row_layout C>
__device__ static inline void mma_ABt(D &d,
const A &a,
const B &b, // notice row and (M, K) instead of col and (K, M)
const C &c) {
KITTENS_CHECK_WARP
static_assert(D::rows == A::rows && D::cols == B::rows); // Check D matches A, B
static_assert(A::cols == B::cols); // Check reduction dim is same
static_assert(D::rows == C::rows && D::cols == C::cols); // Check D matches C
#ifdef KITTENS_HOPPER
static_assert(
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, bf16> &&
std::is_same_v<typename B::T, bf16> && std::is_same_v<typename C::T, float>) ||
(std::is_same_v<typename D::T, half> && std::is_same_v<typename A::T, half> &&
std::is_same_v<typename B::T, half> && std::is_same_v<typename C::T, half>) ||
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, fp8e4m3> &&
std::is_same_v<typename B::T, fp8e4m3> && std::is_same_v<typename C::T, float>)
);
#else
static_assert(
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, bf16> &&
std::is_same_v<typename B::T, bf16> && std::is_same_v<typename C::T, float>) ||
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, half> &&
std::is_same_v<typename B::T, half> && std::is_same_v<typename C::T, float>) ||
(std::is_same_v<typename D::T, half> && std::is_same_v<typename A::T, half> &&
std::is_same_v<typename B::T, half> && std::is_same_v<typename C::T, half>)
);
#endif
#pragma unroll
for(int n = 0; n < D::height; n++) {
#pragma unroll
for(int m = 0; m < D::width; m++) {
mma_ABt_base(
d.tiles[n][m],
a.tiles[n][0],
b.tiles[m][0],
c.tiles[n][m]
);
#pragma unroll
for(int k = 1; k < A::width; k++) {
mma_ABt_base(
d.tiles[n][m],
a.tiles[n][k],
b.tiles[m][k],
d.tiles[n][m]
);
}
}
}
}
/**
* @brief Matrix multiply-accumulate operation with transposed A.
*
* This function performs the matrix multiply-accumulate operation
* using the `hmma16816` instruction.
*
* @tparam N The number of row tiles.
* @tparam K The number of column tiles for the A matrix and row tiles for the B matrix.
* @tparam M The number of column tiles for the B matrix.
* @param[out] d The output rt_fl<N, M, row_layout> accumulator.
* @param[in] a The first input rt_bf<K, N, row_layout> matrix.
* @param[in] b The second input rt_bf<K, M, col_layout> matrix in column-major mode.
* @param[in] c The input rt_fl<N, M, row_layout> accumulator matrix.
*/
template<ducks::rt::row_layout D, ducks::rt::col_layout A, ducks::rt::col_layout B, ducks::rt::row_layout C>
__device__ static inline void mma_AtB(D &d,
const A &a,
const B &b,
const C &c) {
KITTENS_CHECK_WARP
static_assert(D::rows == A::cols && D::cols == B::cols); // Check D matches A, B
static_assert(A::rows == B::rows); // Check reduction dim is same
static_assert(D::rows == C::rows && D::cols == C::cols); // Check D matches C
#ifdef KITTENS_HOPPER
static_assert(
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, bf16> &&
std::is_same_v<typename B::T, bf16> && std::is_same_v<typename C::T, float>) ||
(std::is_same_v<typename D::T, half> && std::is_same_v<typename A::T, half> &&
std::is_same_v<typename B::T, half> && std::is_same_v<typename C::T, half>) ||
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, fp8e4m3> &&
std::is_same_v<typename B::T, fp8e4m3> && std::is_same_v<typename C::T, float>)
);
#else
static_assert(
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, bf16> &&
std::is_same_v<typename B::T, bf16> && std::is_same_v<typename C::T, float>) ||
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, half> &&
std::is_same_v<typename B::T, half> && std::is_same_v<typename C::T, float>) ||
(std::is_same_v<typename D::T, half> && std::is_same_v<typename A::T, half> &&
std::is_same_v<typename B::T, half> && std::is_same_v<typename C::T, half>)
);
#endif
#pragma unroll
for(int n = 0; n < D::height; n++) {
#pragma unroll
for(int m = 0; m < D::width; m++) {
mma_AtB_base(
d.tiles[n][m],
a.tiles[0][n],
b.tiles[0][m],
c.tiles[n][m]
);
#pragma unroll
for(int k = 1; k < A::height; k++) {
mma_AtB_base(
d.tiles[n][m],
a.tiles[k][n],
b.tiles[k][m],
d.tiles[n][m]
);
}
}
}
}
/**
* @brief Matrix multiply-accumulate operation with transposed A and B.
*
* This function performs the matrix multiply-accumulate operation
* using the `hmma16816` instruction.
*
* @tparam N The number of row tiles.
* @tparam K The number of column tiles for the A matrix and row tiles for the B matrix.
* @tparam M The number of column tiles for the B matrix.
* @param[out] d The output rt_fl<N, M, row_layout> accumulator.
* @param[in] a The first input rt_bf<K, N, col_layout> matrix.
* @param[in] b The second input rt_bf<M, K, row_layout> matrix in column-major mode.
* @param[in] c The input rt_fl<N, M, row_layout> accumulator matrix.
*/
template<ducks::rt::row_layout D, ducks::rt::col_layout A, ducks::rt::row_layout B, ducks::rt::row_layout C>
__device__ static inline void mma_AtBt(D &d,
const A &a,
const B &b,
const C &c) {
KITTENS_CHECK_WARP
static_assert(D::rows == A::cols && D::cols == B::rows); // Check D matches A, B
static_assert(A::rows == B::cols); // Check reduction dim is same
static_assert(D::rows == C::rows && D::cols == C::cols); // Check D matches C
#ifdef KITTENS_HOPPER
static_assert(
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, bf16> &&
std::is_same_v<typename B::T, bf16> && std::is_same_v<typename C::T, float>) ||
(std::is_same_v<typename D::T, half> && std::is_same_v<typename A::T, half> &&
std::is_same_v<typename B::T, half> && std::is_same_v<typename C::T, half>) ||
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, fp8e4m3> &&
std::is_same_v<typename B::T, fp8e4m3> && std::is_same_v<typename C::T, float>)
);
#else
static_assert(
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, bf16> &&
std::is_same_v<typename B::T, bf16> && std::is_same_v<typename C::T, float>) ||
(std::is_same_v<typename D::T, float> && std::is_same_v<typename A::T, half> &&
std::is_same_v<typename B::T, half> && std::is_same_v<typename C::T, float>) ||
(std::is_same_v<typename D::T, half> && std::is_same_v<typename A::T, half> &&
std::is_same_v<typename B::T, half> && std::is_same_v<typename C::T, half>)
);
#endif
#pragma unroll
for(int n = 0; n < D::height; n++) {
#pragma unroll
for(int m = 0; m < D::width; m++) {
mma_AtBt_base(
d.tiles[n][m],
a.tiles[0][n],
b.tiles[m][0],
c.tiles[n][m]
);
#pragma unroll
for(int k = 1; k < A::height; k++) {
mma_AtBt_base(
d.tiles[n][m],
a.tiles[k][n],
b.tiles[m][k],
d.tiles[n][m]
);
}
}
}
}
template<int trans_A, int trans_B, ducks::rt::all D, ducks::rt::all A, ducks::rt::all B, ducks::rt::all C>
__device__ static inline void mma(D &d,
const A &a,
const B &b,
const C &c) {
KITTENS_CHECK_WARP
if constexpr(trans_A == transpose::T) {
if constexpr(trans_B == transpose::T) {
mma_AtBt(d, a, b, c);
} else {
mma_AtB(d, a, b, c);
}
} else {
if constexpr(trans_B == transpose::T) {
mma_ABt(d, a, b, c);
} else {
mma_AB(d, a, b, c);
}
}
}
template<int trans_A, int trans_B, ducks::rt::all A, ducks::rt::all B, ducks::rt::all C>
__device__ static inline C mma(const A &a,
const B &b,
const C &c) {
KITTENS_CHECK_WARP
C d;
if constexpr(trans_A == transpose::T) {
if constexpr(trans_B == transpose::T) {
mma_AtBt(d, a, b, c);
} else {
mma_AtB(d, a, b, c);
}
} else {
if constexpr(trans_B == transpose::T) {
mma_ABt(d, a, b, c);
} else {
mma_AB(d, a, b, c);
}
}
return d;
}
// --------------------------------------------------------------------------------------------------------------------
// --------------------------------------------------------------------------------------------------------------------
// -------------------------------------------------- COMPLEX INPUTS --------------------------------------------------
// --------------------------------------------------------------------------------------------------------------------
// --------------------------------------------------------------------------------------------------------------------
/**
* @brief Matrix multiply-accumulate operation for complex tiles
*
* This function calls mma_AB with hf arguments
*
* @tparam N The number of row tiles.
* @tparam K The number of column tiles for the A matrix and row tiles for the B matrix.
* @tparam M The number of column tiles for the B matrix.
* @param[out] d The output rt_cmplx_hf<N, M, row_layout> accumulator.
* @param[in] a The first input rt_cmplx_hf<N, K, row_layout> matrix.
* @param[in] b The second input rt_cmplx_hf<K, M, col_layout> matrix in column-major mode.
* @param[in] c The input rt_cmplx_hf<N, M, row_layout> accumulator matrix.
*/
template<int N, int K, int M>
__device__ static inline void mma_AB(crt_hf<N, M, ducks::rt_layout::row> &d,
const crt_hf<N, K, ducks::rt_layout::row> &a,
const crt_hf<K, M, ducks::rt_layout::col> &b,
const crt_hf<N, M, ducks::rt_layout::row> &c) {
KITTENS_CHECK_WARP
// Copy data from input accumulate register into output
::kittens::group<1>::copy(d.real, c.real);
::kittens::group<1>::copy(d.imag, c.imag);
// Negative on B matrix so we can use single accum register
rt_hf<N, K, ducks::rt_layout::row> tmp;
// Hex value for -1 in float16
constexpr half factor = std::bit_cast<__half>(uint16_t(0xFB80));
::kittens::group<1>::mul(tmp, a.imag, factor);
mma_AB(d.real, a.real, b.real, d.real);
mma_AB(d.real, tmp, b.imag, d.real);
mma_AB(d.imag, a.real, b.imag, d.imag);
mma_AB(d.imag, a.imag, b.real, d.imag);
}
/**
* @brief Matrix multiply-accumulate operation for complex tiles
*
* This function calls mma_AB with bf16 arguments
*
* @tparam N The number of row tiles.
* @tparam K The number of column tiles for the A matrix and row tiles for the B matrix.
* @tparam M The number of column tiles for the B matrix.
* @param[out] d The output rt_cmplx_fl<N, M, row_layout> accumulator.
* @param[in] a The first input rt_cmplx_bf<N, K, row_layout> matrix.
* @param[in] b The second input rt_cmplx_bf<K, M, col_layout> matrix in column-major mode.
* @param[in] c The input rt_cmplx_fl<N, M, row_layout> accumulator matrix.
*/
template<int N, int K, int M>
__device__ static inline void mma_AB(crt_fl<N, M, ducks::rt_layout::row> &d,
const crt_bf<N, K, ducks::rt_layout::row> &a,
const crt_bf<K, M, ducks::rt_layout::col> &b,
const crt_fl<N, M, ducks::rt_layout::row> &c) {
KITTENS_CHECK_WARP
// Copy data from input accumulate register into output
::kittens::group<1>::copy(d.real, c.real);
::kittens::group<1>::copy(d.imag, c.imag);
// Negative on B matrix so we can use single accum register
kittens::rt_bf<N, K, ducks::rt_layout::row> tmp;
// Hex value for -1 in bf16
constexpr bf16 factor = std::bit_cast<__nv_bfloat16>(uint16_t(0xBF80));
::kittens::group<1>::mul(tmp, a.imag, factor);
mma_AB(d.real, a.real, b.real, d.real);
mma_AB(d.real, tmp, b.imag, d.real);
mma_AB(d.imag, a.real, b.imag, d.imag);
mma_AB(d.imag, a.imag, b.real, d.imag);
}
@@ -1,334 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 112, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 112, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %61, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n112k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \
"{%56, %57, %58, %59}, " \
"%60, " \
"p, 1, %63, %62;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y),
"+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y),
"+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y),
"+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y),
"+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y),
"+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y),
"+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y),
"+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y),
"+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %61, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \
"{%56, %57, %58, %59}, " \
"%60, " \
"p, 1, %63, %62;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y),
"+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y),
"+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y),
"+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y),
"+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y),
"+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y),
"+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y),
"+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y),
"+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %33, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n112k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27}, " \
"{%28, %29, %30, %31}, " \
"%32, " \
"p, 1, %35, %34;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 112, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %58, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n112k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \
"%56, " \
"%57, " \
"p, 1, %61, %59, %60;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y),
"+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y),
"+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y),
"+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y),
"+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y),
"+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y),
"+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y),
"+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y),
"+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %58, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n112k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \
"%56, " \
"%57, " \
"p, 1, %61, %59, %60;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y),
"+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y),
"+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y),
"+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y),
"+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y),
"+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y),
"+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y),
"+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y),
"+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %30, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n112k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27}, " \
"%28, " \
"%29, " \
"p, 1, %33, %31, %32;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
}
};
@@ -1,813 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 128, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 128, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e5m2>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %69, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \
"{%64, %65, %66, %67}, " \
"%68, " \
"p, 1, %71, %70;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y),
"+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y),
"+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y),
"+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y),
"+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y),
"+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y),
"+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y),
"+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y),
"+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y),
"+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y),
"+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y),
"+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y),
"+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %69, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \
"{%64, %65, %66, %67}, " \
"%68, " \
"p, 1, %71, %70;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y),
"+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y),
"+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y),
"+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y),
"+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y),
"+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y),
"+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y),
"+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y),
"+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y),
"+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y),
"+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y),
"+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y),
"+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %37, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"{%32, %33, %34, %35}, " \
"%36, " \
"p, 1, %39, %38;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %69, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k32.f32.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \
"{%64, %65, %66, %67}, " \
"%68, " \
"p, 1, %70;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y),
"+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y),
"+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y),
"+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y),
"+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y),
"+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y),
"+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y),
"+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y),
"+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y),
"+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y),
"+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y),
"+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y),
"+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %69, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k32.f32.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \
"{%64, %65, %66, %67}, " \
"%68, " \
"p, 1, %70;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y),
"+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y),
"+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y),
"+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y),
"+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y),
"+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y),
"+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y),
"+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y),
"+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y),
"+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y),
"+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y),
"+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y),
"+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %37, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k32.f16.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"{%32, %33, %34, %35}, " \
"%36, " \
"p, 1, %38;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %37, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k32.f16.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"{%32, %33, %34, %35}, " \
"%36, " \
"p, 1, %38;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 128, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e5m2>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %66, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \
"%64, " \
"%65, " \
"p, 1, %69, %67, %68;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y),
"+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y),
"+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y),
"+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y),
"+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y),
"+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y),
"+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y),
"+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y),
"+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y),
"+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y),
"+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y),
"+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y),
"+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %66, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \
"%64, " \
"%65, " \
"p, 1, %69, %67, %68;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y),
"+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y),
"+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y),
"+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y),
"+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y),
"+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y),
"+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y),
"+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y),
"+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y),
"+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y),
"+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y),
"+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y),
"+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %34, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"%32, " \
"%33, " \
"p, 1, %37, %35, %36;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %66, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k32.f32.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \
"%64, " \
"%65, " \
"p, 1, %67;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y),
"+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y),
"+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y),
"+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y),
"+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y),
"+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y),
"+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y),
"+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y),
"+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y),
"+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y),
"+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y),
"+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y),
"+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %66, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k32.f32.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, " \
"%64, " \
"%65, " \
"p, 1, %67;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y),
"+f"(dst.tiles[0][5].data[0].x), "+f"(dst.tiles[0][5].data[0].y),
"+f"(dst.tiles[0][5].data[1].x), "+f"(dst.tiles[0][5].data[1].y),
"+f"(dst.tiles[0][5].data[2].x), "+f"(dst.tiles[0][5].data[2].y),
"+f"(dst.tiles[0][5].data[3].x), "+f"(dst.tiles[0][5].data[3].y),
"+f"(dst.tiles[0][6].data[0].x), "+f"(dst.tiles[0][6].data[0].y),
"+f"(dst.tiles[0][6].data[1].x), "+f"(dst.tiles[0][6].data[1].y),
"+f"(dst.tiles[0][6].data[2].x), "+f"(dst.tiles[0][6].data[2].y),
"+f"(dst.tiles[0][6].data[3].x), "+f"(dst.tiles[0][6].data[3].y),
"+f"(dst.tiles[0][7].data[0].x), "+f"(dst.tiles[0][7].data[0].y),
"+f"(dst.tiles[0][7].data[1].x), "+f"(dst.tiles[0][7].data[1].y),
"+f"(dst.tiles[0][7].data[2].x), "+f"(dst.tiles[0][7].data[2].y),
"+f"(dst.tiles[0][7].data[3].x), "+f"(dst.tiles[0][7].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %34, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k32.f16.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"%32, " \
"%33, " \
"p, 1, %35;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %34, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n128k32.f16.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"%32, " \
"%33, " \
"p, 1, %35;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][7].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
}
};
@@ -1,382 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 144, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 144, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %77, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n144k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71}, " \
"{%72, %73, %74, %75}, " \
"%76, " \
"p, 1, %79, %78;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %77, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71}, " \
"{%72, %73, %74, %75}, " \
"%76, " \
"p, 1, %79, %78;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %41, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n144k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35}, " \
"{%36, %37, %38, %39}, " \
"%40, " \
"p, 1, %43, %42;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 144, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %74, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n144k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71}, " \
"%72, " \
"%73, " \
"p, 1, %77, %75, %76;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %74, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n144k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71}, " \
"%72, " \
"%73, " \
"p, 1, %77, %75, %76;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %38, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n144k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35}, " \
"%36, " \
"%37, " \
"p, 1, %41, %39, %40;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
}
};
@@ -1,190 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 16, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 16, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %13, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7}, " \
"{%8, %9, %10, %11}, " \
"%12, " \
"p, 1, %15, %14;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %13, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7}, " \
"{%8, %9, %10, %11}, " \
"%12, " \
"p, 1, %15, %14;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %9, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n16k16.f16.f16.f16 " \
"{%0, %1, %2, %3}, " \
"{%4, %5, %6, %7}, " \
"%8, " \
"p, 1, %11, %10;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 16, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %10, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7}, " \
"%8, " \
"%9, " \
"p, 1, %13, %11, %12;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %10, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7}, " \
"%8, " \
"%9, " \
"p, 1, %13, %11, %12;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %6, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n16k16.f16.f16.f16 " \
"{%0, %1, %2, %3}, " \
"%4, " \
"%5, " \
"p, 1, %9, %7, %8;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
}
};
@@ -1,666 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 160, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 160, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %85, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n160k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \
"{%80, %81, %82, %83}, " \
"%84, " \
"p, 1, %87, %86;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %85, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \
"{%80, %81, %82, %83}, " \
"%84, " \
"p, 1, %87, %86;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %45, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n160k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \
"{%40, %41, %42, %43}, " \
"%44, " \
"p, 1, %47, %46;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %85, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n160k32.f32.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \
"{%80, %81, %82, %83}, " \
"%84, " \
"p, 1, %86;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %85, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n160k32.f32.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \
"{%80, %81, %82, %83}, " \
"%84, " \
"p, 1, %86;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 160, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %82, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n160k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \
"%80, " \
"%81, " \
"p, 1, %85, %83, %84;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %82, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n160k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \
"%80, " \
"%81, " \
"p, 1, %85, %83, %84;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %42, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n160k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \
"%40, " \
"%41, " \
"p, 1, %45, %43, %44;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %82, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n160k32.f32.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \
"%80, " \
"%81, " \
"p, 1, %83;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %82, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n160k32.f32.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79}, " \
"%80, " \
"%81, " \
"p, 1, %83;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
}
};
@@ -1,430 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 176, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 176, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %93, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n176k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87}, " \
"{%88, %89, %90, %91}, " \
"%92, " \
"p, 1, %95, %94;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %93, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87}, " \
"{%88, %89, %90, %91}, " \
"%92, " \
"p, 1, %95, %94;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %49, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n176k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43}, " \
"{%44, %45, %46, %47}, " \
"%48, " \
"p, 1, %51, %50;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 176, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %90, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n176k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87}, " \
"%88, " \
"%89, " \
"p, 1, %93, %91, %92;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %90, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n176k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87}, " \
"%88, " \
"%89, " \
"p, 1, %93, %91, %92;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %46, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n176k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43}, " \
"%44, " \
"%45, " \
"p, 1, %49, %47, %48;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
}
};
@@ -1,674 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 192, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 192, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %101, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \
"{%96, %97, %98, %99}, " \
"%100, " \
"p, 1, %103, %102;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %101, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \
"{%96, %97, %98, %99}, " \
"%100, " \
"p, 1, %103, %102;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %53, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n192k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \
"{%48, %49, %50, %51}, " \
"%52, " \
"p, 1, %55, %54;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %101, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n192k32.f32.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \
"{%96, %97, %98, %99}, " \
"%100, " \
"p, 1, %102;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %101, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n192k32.f32.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \
"{%96, %97, %98, %99}, " \
"%100, " \
"p, 1, %102;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 192, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %98, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \
"%96, " \
"%97, " \
"p, 1, %101, %99, %100;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %98, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \
"%96, " \
"%97, " \
"p, 1, %101, %99, %100;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %50, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n192k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47}, " \
"%48, " \
"%49, " \
"p, 1, %53, %51, %52;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %98, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n192k32.f32.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95}, " \
"%96, " \
"%97, " \
"p, 1, %99;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
}
};
@@ -1,478 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 208, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 208, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %109, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n208k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103}, " \
"{%104, %105, %106, %107}, " \
"%108, " \
"p, 1, %111, %110;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %109, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103}, " \
"{%104, %105, %106, %107}, " \
"%108, " \
"p, 1, %111, %110;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %57, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n208k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51}, " \
"{%52, %53, %54, %55}, " \
"%56, " \
"p, 1, %59, %58;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 208, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %106, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n208k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103}, " \
"%104, " \
"%105, " \
"p, 1, %109, %107, %108;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %106, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n208k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103}, " \
"%104, " \
"%105, " \
"p, 1, %109, %107, %108;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %54, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n208k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51}, " \
"%52, " \
"%53, " \
"p, 1, %57, %55, %56;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
}
};
@@ -1,826 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 224, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 224, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %117, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n224k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \
"{%112, %113, %114, %115}, " \
"%116, " \
"p, 1, %119, %118;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y),
"+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y),
"+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y),
"+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y),
"+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %117, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \
"{%112, %113, %114, %115}, " \
"%116, " \
"p, 1, %119, %118;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y),
"+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y),
"+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y),
"+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y),
"+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %61, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n224k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \
"{%56, %57, %58, %59}, " \
"%60, " \
"p, 1, %63, %62;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %117, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n224k32.f32.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \
"{%112, %113, %114, %115}, " \
"%116, " \
"p, 1, %118;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y),
"+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y),
"+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y),
"+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y),
"+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %117, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n224k32.f32.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \
"{%112, %113, %114, %115}, " \
"%116, " \
"p, 1, %118;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y),
"+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y),
"+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y),
"+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y),
"+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 224, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %114, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n224k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \
"%112, " \
"%113, " \
"p, 1, %117, %115, %116;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y),
"+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y),
"+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y),
"+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y),
"+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %114, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n224k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \
"%112, " \
"%113, " \
"p, 1, %117, %115, %116;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y),
"+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y),
"+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y),
"+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y),
"+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %58, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n224k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55}, " \
"%56, " \
"%57, " \
"p, 1, %61, %59, %60;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %114, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n224k32.f32.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \
"%112, " \
"%113, " \
"p, 1, %117, %115, %116;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y),
"+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y),
"+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y),
"+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y),
"+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %114, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n224k32.f32.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111}, " \
"%112, " \
"%113, " \
"p, 1, %117, %115, %116;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y),
"+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y),
"+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y),
"+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y),
"+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
}
};
@@ -1,526 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 240, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 240, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %125, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n240k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119}, " \
"{%120, %121, %122, %123}, " \
"%124, " \
"p, 1, %127, %126;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y),
"+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y),
"+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y),
"+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y),
"+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y),
"+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y),
"+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y),
"+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y),
"+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %125, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119}, " \
"{%120, %121, %122, %123}, " \
"%124, " \
"p, 1, %127, %126;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y),
"+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y),
"+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y),
"+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y),
"+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y),
"+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y),
"+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y),
"+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y),
"+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %65, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n240k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59}, " \
"{%60, %61, %62, %63}, " \
"%64, " \
"p, 1, %67, %66;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][14].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][14].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][14].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][14].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 240, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %122, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n240k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119}, " \
"%120, " \
"%121, " \
"p, 1, %125, %123, %124;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y),
"+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y),
"+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y),
"+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y),
"+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y),
"+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y),
"+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y),
"+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y),
"+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %122, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n240k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119}, " \
"%120, " \
"%121, " \
"p, 1, %125, %123, %124;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][ 0].data[0].x), "+f"(dst.tiles[0][ 0].data[0].y),
"+f"(dst.tiles[0][ 0].data[1].x), "+f"(dst.tiles[0][ 0].data[1].y),
"+f"(dst.tiles[0][ 0].data[2].x), "+f"(dst.tiles[0][ 0].data[2].y),
"+f"(dst.tiles[0][ 0].data[3].x), "+f"(dst.tiles[0][ 0].data[3].y),
"+f"(dst.tiles[0][ 1].data[0].x), "+f"(dst.tiles[0][ 1].data[0].y),
"+f"(dst.tiles[0][ 1].data[1].x), "+f"(dst.tiles[0][ 1].data[1].y),
"+f"(dst.tiles[0][ 1].data[2].x), "+f"(dst.tiles[0][ 1].data[2].y),
"+f"(dst.tiles[0][ 1].data[3].x), "+f"(dst.tiles[0][ 1].data[3].y),
"+f"(dst.tiles[0][ 2].data[0].x), "+f"(dst.tiles[0][ 2].data[0].y),
"+f"(dst.tiles[0][ 2].data[1].x), "+f"(dst.tiles[0][ 2].data[1].y),
"+f"(dst.tiles[0][ 2].data[2].x), "+f"(dst.tiles[0][ 2].data[2].y),
"+f"(dst.tiles[0][ 2].data[3].x), "+f"(dst.tiles[0][ 2].data[3].y),
"+f"(dst.tiles[0][ 3].data[0].x), "+f"(dst.tiles[0][ 3].data[0].y),
"+f"(dst.tiles[0][ 3].data[1].x), "+f"(dst.tiles[0][ 3].data[1].y),
"+f"(dst.tiles[0][ 3].data[2].x), "+f"(dst.tiles[0][ 3].data[2].y),
"+f"(dst.tiles[0][ 3].data[3].x), "+f"(dst.tiles[0][ 3].data[3].y),
"+f"(dst.tiles[0][ 4].data[0].x), "+f"(dst.tiles[0][ 4].data[0].y),
"+f"(dst.tiles[0][ 4].data[1].x), "+f"(dst.tiles[0][ 4].data[1].y),
"+f"(dst.tiles[0][ 4].data[2].x), "+f"(dst.tiles[0][ 4].data[2].y),
"+f"(dst.tiles[0][ 4].data[3].x), "+f"(dst.tiles[0][ 4].data[3].y),
"+f"(dst.tiles[0][ 5].data[0].x), "+f"(dst.tiles[0][ 5].data[0].y),
"+f"(dst.tiles[0][ 5].data[1].x), "+f"(dst.tiles[0][ 5].data[1].y),
"+f"(dst.tiles[0][ 5].data[2].x), "+f"(dst.tiles[0][ 5].data[2].y),
"+f"(dst.tiles[0][ 5].data[3].x), "+f"(dst.tiles[0][ 5].data[3].y),
"+f"(dst.tiles[0][ 6].data[0].x), "+f"(dst.tiles[0][ 6].data[0].y),
"+f"(dst.tiles[0][ 6].data[1].x), "+f"(dst.tiles[0][ 6].data[1].y),
"+f"(dst.tiles[0][ 6].data[2].x), "+f"(dst.tiles[0][ 6].data[2].y),
"+f"(dst.tiles[0][ 6].data[3].x), "+f"(dst.tiles[0][ 6].data[3].y),
"+f"(dst.tiles[0][ 7].data[0].x), "+f"(dst.tiles[0][ 7].data[0].y),
"+f"(dst.tiles[0][ 7].data[1].x), "+f"(dst.tiles[0][ 7].data[1].y),
"+f"(dst.tiles[0][ 7].data[2].x), "+f"(dst.tiles[0][ 7].data[2].y),
"+f"(dst.tiles[0][ 7].data[3].x), "+f"(dst.tiles[0][ 7].data[3].y),
"+f"(dst.tiles[0][ 8].data[0].x), "+f"(dst.tiles[0][ 8].data[0].y),
"+f"(dst.tiles[0][ 8].data[1].x), "+f"(dst.tiles[0][ 8].data[1].y),
"+f"(dst.tiles[0][ 8].data[2].x), "+f"(dst.tiles[0][ 8].data[2].y),
"+f"(dst.tiles[0][ 8].data[3].x), "+f"(dst.tiles[0][ 8].data[3].y),
"+f"(dst.tiles[0][ 9].data[0].x), "+f"(dst.tiles[0][ 9].data[0].y),
"+f"(dst.tiles[0][ 9].data[1].x), "+f"(dst.tiles[0][ 9].data[1].y),
"+f"(dst.tiles[0][ 9].data[2].x), "+f"(dst.tiles[0][ 9].data[2].y),
"+f"(dst.tiles[0][ 9].data[3].x), "+f"(dst.tiles[0][ 9].data[3].y),
"+f"(dst.tiles[0][10].data[0].x), "+f"(dst.tiles[0][10].data[0].y),
"+f"(dst.tiles[0][10].data[1].x), "+f"(dst.tiles[0][10].data[1].y),
"+f"(dst.tiles[0][10].data[2].x), "+f"(dst.tiles[0][10].data[2].y),
"+f"(dst.tiles[0][10].data[3].x), "+f"(dst.tiles[0][10].data[3].y),
"+f"(dst.tiles[0][11].data[0].x), "+f"(dst.tiles[0][11].data[0].y),
"+f"(dst.tiles[0][11].data[1].x), "+f"(dst.tiles[0][11].data[1].y),
"+f"(dst.tiles[0][11].data[2].x), "+f"(dst.tiles[0][11].data[2].y),
"+f"(dst.tiles[0][11].data[3].x), "+f"(dst.tiles[0][11].data[3].y),
"+f"(dst.tiles[0][12].data[0].x), "+f"(dst.tiles[0][12].data[0].y),
"+f"(dst.tiles[0][12].data[1].x), "+f"(dst.tiles[0][12].data[1].y),
"+f"(dst.tiles[0][12].data[2].x), "+f"(dst.tiles[0][12].data[2].y),
"+f"(dst.tiles[0][12].data[3].x), "+f"(dst.tiles[0][12].data[3].y),
"+f"(dst.tiles[0][13].data[0].x), "+f"(dst.tiles[0][13].data[0].y),
"+f"(dst.tiles[0][13].data[1].x), "+f"(dst.tiles[0][13].data[1].y),
"+f"(dst.tiles[0][13].data[2].x), "+f"(dst.tiles[0][13].data[2].y),
"+f"(dst.tiles[0][13].data[3].x), "+f"(dst.tiles[0][13].data[3].y),
"+f"(dst.tiles[0][14].data[0].x), "+f"(dst.tiles[0][14].data[0].y),
"+f"(dst.tiles[0][14].data[1].x), "+f"(dst.tiles[0][14].data[1].y),
"+f"(dst.tiles[0][14].data[2].x), "+f"(dst.tiles[0][14].data[2].y),
"+f"(dst.tiles[0][14].data[3].x), "+f"(dst.tiles[0][14].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %62, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n240k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59}, " \
"%60, " \
"%61, " \
"p, 1, %65, %63, %64;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][ 0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 4].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 5].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 6].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 7].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 8].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][ 9].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][10].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][11].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][12].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][13].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][14].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][14].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][14].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][14].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
}
};
File diff suppressed because it is too large Load Diff
@@ -1,446 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 32, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 32, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e5m2>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %21, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"{%16, %17, %18, %19}, " \
"%20, " \
"p, 1, %23, %22;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %21, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"{%16, %17, %18, %19}, " \
"%20, " \
"p, 1, %23, %22;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %13, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7}, " \
"{%8, %9, %10, %11}, " \
"%12, " \
"p, 1, %15, %14;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %21, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k32.f32.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"{%16, %17, %18, %19}, " \
"%20, " \
"p, 1, %22;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %21, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k32.f32.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"{%16, %17, %18, %19}, " \
"%20, " \
"p, 1, %22;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %13, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k32.f16.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7}, " \
"{%8, %9, %10, %11}, " \
"%12, " \
"p, 1, %14;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 32, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e5m2>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %18, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"%16, " \
"%17, " \
"p, 1, %21, %19, %20;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %18, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"%16, " \
"%17, " \
"p, 1, %21, %19, %20;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %10, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7}, " \
"%8, " \
"%9, " \
"p, 1, %13, %11, %12;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %18, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k32.f32.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"%16, " \
"%17, " \
"p, 1, %19;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %18, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k32.f32.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"%16, " \
"%17, " \
"p, 1, %19;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %10, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k32.f16.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7}, " \
"%8, " \
"%9, " \
"p, 1, %11;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %10, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n32k32.f16.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7}, " \
"%8, " \
"%9, " \
"p, 1, %11;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
}
};
@@ -1,238 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 48, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 48, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %29, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n48k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23}, " \
"{%24, %25, %26, %27}, " \
"%28, " \
"p, 1, %31, %30;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %29, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23}, " \
"{%24, %25, %26, %27}, " \
"%28, " \
"p, 1, %31, %30;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %17, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n48k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11}, " \
"{%12, %13, %14, %15}, " \
"%16, " \
"p, 1, %19, %18;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 48, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %26, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n48k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23}, " \
"%24, " \
"%25, " \
"p, 1, %29, %27, %28;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %26, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n48k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23}, " \
"%24, " \
"%25, " \
"p, 1, %29, %27, %28;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %14, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n48k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11}, " \
"%12, " \
"%13, " \
"p, 1, %17, %15, %16;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
}
};
@@ -1,587 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 64, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 64, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e5m2>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %37, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"{%32, %33, %34, %35}, " \
"%36, " \
"p, 1, %39, %38;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %37, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"{%32, %33, %34, %35}, " \
"%36, " \
"p, 1, %39, %38;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %21, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"{%16, %17, %18, %19}, " \
"%20, " \
"p, 1, %23, %22;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %37, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k32.f32.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"{%32, %33, %34, %35}, " \
"%36, " \
"p, 1, %38;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %37, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k32.f32.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"{%32, %33, %34, %35}, " \
"%36, " \
"p, 1, %38;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %21, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k32.f16.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"{%16, %17, %18, %19}, " \
"%20, " \
"p, 1, %22;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %21, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k32.f16.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"{%16, %17, %18, %19}, " \
"%20, " \
"p, 1, %22;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d),
// "n"(trans_b),
"n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 64, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e4m3>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e5m2>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %34, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"%32, " \
"%33, " \
"p, 1, %37, %35, %36;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %34, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"%32, " \
"%33, " \
"p, 1, %37, %35, %36;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %18, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"%16, " \
"%17, " \
"p, 1, %21, %19, %20;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %34, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k32.f32.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"%32, " \
"%33, " \
"p, 1, %35;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b), // transpose is not supported for FP8
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %34, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k32.f32.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, " \
"%32, " \
"%33, " \
"p, 1, %35;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b), // transpose is not supported for FP8
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e4m3>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %18, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k32.f16.e4m3.e4m3 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"%16, " \
"%17, " \
"p, 1, %19;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
// ----- FP8,FP8 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, fp8e5m2>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %18, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n64k32.f16.e5m2.e5m2 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, " \
"%16, " \
"%17, " \
"p, 1, %19;\n" \
"}\n"
// a_mat descriptor, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a, imm-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
// "n"(trans_a),
// "n"(trans_b),
"n"(scale_b)
);
}
}
};
@@ -1,286 +0,0 @@
template<typename T_D, typename T_AB, int trans_a, int trans_b>
struct base<T_D, T_AB, 80, trans_a, trans_b> {
template<int scale_b=1> __device__ static inline void rt_st(
rt<T_D, 16, 80, ducks::rt_layout::row> &dst,
const rt_base<T_AB, ducks::rt_layout::row> & a_rt,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %45, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n80k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \
"{%40, %41, %42, %43}, " \
"%44, " \
"p, 1, %47, %46;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %45, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \
"{%40, %41, %42, %43}, " \
"%44, " \
"p, 1, %47, %46;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y)
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %25, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n80k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19}, " \
"{%20, %21, %22, %23}, " \
"%24, " \
"p, 1, %27, %26;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[3])
: "r"(*(uint32_t*)&a_rt.data[0]), "r"(*(uint32_t*)&a_rt.data[1]),
"r"(*(uint32_t*)&a_rt.data[2]), "r"(*(uint32_t*)&a_rt.data[3]),
"l"(b_st_desc), "r"(scale_d), "n"(trans_b), "n"(scale_b)
);
}
}
template<int scale_b=1> __device__ static inline void st_st(
rt<T_D, 16, 80, ducks::rt_layout::row> &dst,
const uint64_t a_st_desc,
const uint64_t b_st_desc,
int scale_d = 1
) {
static_assert(
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) ||
(std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) ||
(std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>),
"Invalid type combination for WGMMA."
);
static_assert(scale_b==1 || scale_b==-1, "Invalid scale B (invert) option");
// ----- BF16,BF16 -> FP32 ----- //
if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, bf16>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %42, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n80k16.f32.bf16.bf16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \
"%40, " \
"%41, " \
"p, 1, %45, %43, %44;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP32 ----- //
else if constexpr (std::is_same_v<T_D, float> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %42, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n80k16.f32.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39}, " \
"%40, " \
"%41, " \
"p, 1, %45, %43, %44;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+f"(dst.tiles[0][0].data[0].x), "+f"(dst.tiles[0][0].data[0].y),
"+f"(dst.tiles[0][0].data[1].x), "+f"(dst.tiles[0][0].data[1].y),
"+f"(dst.tiles[0][0].data[2].x), "+f"(dst.tiles[0][0].data[2].y),
"+f"(dst.tiles[0][0].data[3].x), "+f"(dst.tiles[0][0].data[3].y),
"+f"(dst.tiles[0][1].data[0].x), "+f"(dst.tiles[0][1].data[0].y),
"+f"(dst.tiles[0][1].data[1].x), "+f"(dst.tiles[0][1].data[1].y),
"+f"(dst.tiles[0][1].data[2].x), "+f"(dst.tiles[0][1].data[2].y),
"+f"(dst.tiles[0][1].data[3].x), "+f"(dst.tiles[0][1].data[3].y),
"+f"(dst.tiles[0][2].data[0].x), "+f"(dst.tiles[0][2].data[0].y),
"+f"(dst.tiles[0][2].data[1].x), "+f"(dst.tiles[0][2].data[1].y),
"+f"(dst.tiles[0][2].data[2].x), "+f"(dst.tiles[0][2].data[2].y),
"+f"(dst.tiles[0][2].data[3].x), "+f"(dst.tiles[0][2].data[3].y),
"+f"(dst.tiles[0][3].data[0].x), "+f"(dst.tiles[0][3].data[0].y),
"+f"(dst.tiles[0][3].data[1].x), "+f"(dst.tiles[0][3].data[1].y),
"+f"(dst.tiles[0][3].data[2].x), "+f"(dst.tiles[0][3].data[2].y),
"+f"(dst.tiles[0][3].data[3].x), "+f"(dst.tiles[0][3].data[3].y),
"+f"(dst.tiles[0][4].data[0].x), "+f"(dst.tiles[0][4].data[0].y),
"+f"(dst.tiles[0][4].data[1].x), "+f"(dst.tiles[0][4].data[1].y),
"+f"(dst.tiles[0][4].data[2].x), "+f"(dst.tiles[0][4].data[2].y),
"+f"(dst.tiles[0][4].data[3].x), "+f"(dst.tiles[0][4].data[3].y)
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
// ----- FP16,FP16 -> FP16 ----- //
else if constexpr (std::is_same_v<T_D, half> && std::is_same_v<T_AB, half>) {
asm volatile (
"{\n"
".reg .pred p;\n" \
"setp.ne.b32 p, %22, 0;\n" \
"wgmma.mma_async.sync.aligned.m64n80k16.f16.f16.f16 " \
"{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19}, " \
"%20, " \
"%21, " \
"p, 1, %25, %23, %24;\n" \
"}\n"
// a_regs, b_mat descriptor, scale-d, imm-scale-a, imm-scale-b, im-trans-a im-trans-b
: "+r"(*(uint32_t*)&dst.tiles[0][0].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][0].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][1].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][2].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][3].data[3]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[0]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[1]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[2]),
"+r"(*(uint32_t*)&dst.tiles[0][4].data[3])
: "l"(a_st_desc),
"l"(b_st_desc),
"r"(scale_d),
"n"(trans_a),
"n"(trans_b),
"n"(scale_b)
);
}
}
};

Some files were not shown because too many files have changed in this diff Show More