forked from tinygrad/tinygrad
Compare commits
247
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f082cbcb36 | ||
|
|
5ad62f130d | ||
|
|
f129d75ee5 | ||
|
|
51f3a5cbb4 | ||
|
|
1d7a8b33c1 | ||
|
|
3fae886aa9 | ||
|
|
514d2a0774 | ||
|
|
7b48f3cc45 | ||
|
|
3f44ef699f | ||
|
|
fa23f37e33 | ||
|
|
284db26a12 | ||
|
|
0a0cb0b9e8 | ||
|
|
a5484b767e | ||
|
|
b4509fba31 | ||
|
|
0f25b4b289 | ||
|
|
f664bcc8bd | ||
|
|
1af05dae77 | ||
|
|
76e8a3250c | ||
|
|
0c015a24fe | ||
|
|
a1881b0c17 | ||
|
|
1b1978b9c0 | ||
|
|
c1e85f699c | ||
|
|
1823a5043f | ||
|
|
46e8ea15c1 | ||
|
|
1216fff781 | ||
|
|
6ad9a688ed | ||
|
|
74b04f7dca | ||
|
|
69857d0ab0 | ||
|
|
a976ace404 | ||
|
|
4b60121498 | ||
|
|
b5f31d7505 | ||
|
|
865d5796f8 | ||
|
|
e74be4a140 | ||
|
|
394dc24110 | ||
|
|
9f2b69b870 | ||
|
|
0b534f71c2 | ||
|
|
b087663c35 | ||
|
|
940a8d5ba9 | ||
|
|
d290e77a5b | ||
|
|
23d310bcc1 | ||
|
|
1e8945a28c | ||
|
|
c7849ac593 | ||
|
|
0f82d92b9d | ||
|
|
4c63f7e786 | ||
|
|
0047bcc535 | ||
|
|
f203d8b221 | ||
|
|
a6dd5a224b | ||
|
|
bf99de7b1e | ||
|
|
9cd365c12e | ||
|
|
16a65b4fd0 | ||
|
|
2d24af888b | ||
|
|
1b58ef0d60 | ||
|
|
17d36d0952 | ||
|
|
7b3912d8e4 | ||
|
|
98163832e4 | ||
|
|
37beef6de3 | ||
|
|
f21851b099 | ||
|
|
ec177c80c2 | ||
|
|
13a25b2e67 | ||
|
|
5d9035f5a6 | ||
|
|
0f804c9a83 | ||
|
|
0eee93f0c0 | ||
|
|
583553f467 | ||
|
|
6fc6b51b59 | ||
|
|
d1c868f990 | ||
|
|
2fcd55583f | ||
|
|
8b48e19ce2 | ||
|
|
3770dd9d80 | ||
|
|
5b649616ff | ||
|
|
9a64fc0d28 | ||
|
|
3e0e0290ce | ||
|
|
2f8ac77c25 | ||
|
|
89bed28716 | ||
|
|
74ee305948 | ||
|
|
f198a9e1ba | ||
|
|
ac3d457d5e | ||
|
|
60e52fbe36 | ||
|
|
6c95b1f39d | ||
|
|
6ba8bf282f | ||
|
|
689ab9151b | ||
|
|
adc8c3b28f | ||
|
|
154d114364 | ||
|
|
fe96c8d345 | ||
|
|
f205352cd7 | ||
|
|
90b1c0dd96 | ||
|
|
42748ccb92 | ||
|
|
05e91a248d | ||
|
|
714500edfd | ||
|
|
57ad46c6e4 | ||
|
|
e02da8f5ac | ||
|
|
0662946fac | ||
|
|
da52006bde | ||
|
|
1c1b4d14e9 | ||
|
|
4204edc60b | ||
|
|
8def8145e4 | ||
|
|
4c9a930de2 | ||
|
|
26247573e1 | ||
|
|
f2eb92948d | ||
|
|
a128fa0f8a | ||
|
|
969a1b35ca | ||
|
|
9ef319f349 | ||
|
|
080b26e7d7 | ||
|
|
44558a37f7 | ||
|
|
2c397eb2a2 | ||
|
|
a83f219253 | ||
|
|
a95159d579 | ||
|
|
9cf5e66899 | ||
|
|
b4a4817c9c | ||
|
|
de1d562b69 | ||
|
|
c9ef5d8fe5 | ||
|
|
e8c595c29e | ||
|
|
360980f1a3 | ||
|
|
109c63b904 | ||
|
|
7129419500 | ||
|
|
4ff7f20b9d | ||
|
|
86c5c969ea | ||
|
|
6a56d3c859 | ||
|
|
ab6b0d3a21 | ||
|
|
2a7310ab59 | ||
|
|
73b25bf47d | ||
|
|
2a0caa09c2 | ||
|
|
881709cd33 | ||
|
|
39aae679e4 | ||
|
|
af935e7d32 | ||
|
|
f522e83a02 | ||
|
|
d95d018bb5 | ||
|
|
05275c9ec3 | ||
|
|
8e508a9927 | ||
|
|
3a480b858f | ||
|
|
32d69d07d7 | ||
|
|
d55d829635 | ||
|
|
c38f6ce140 | ||
|
|
c2689c505e | ||
|
|
cdfa0f29fd | ||
|
|
baf3b60cfb | ||
|
|
9513f025c5 | ||
|
|
b899392f30 | ||
|
|
7ae6898e31 | ||
|
|
3291e00df7 | ||
|
|
9d2f2b8e34 | ||
|
|
9915bcf2b4 | ||
|
|
76c87d81b3 | ||
|
|
fd2e4f2353 | ||
|
|
29469577e8 | ||
|
|
a982480512 | ||
|
|
e01a3eb59a | ||
|
|
cf925d1ac5 | ||
|
|
b252f890da | ||
|
|
292cb6ae26 | ||
|
|
250cb10e8f | ||
|
|
ed90de6583 | ||
|
|
29f0886395 | ||
|
|
b98f1881ef | ||
|
|
6f1cf717de | ||
|
|
0104b16b9b | ||
|
|
f5eb46a3d9 | ||
|
|
8b2e0930d7 | ||
|
|
74411984fc | ||
|
|
d2cd269e28 | ||
|
|
17cec8d645 | ||
|
|
476a2a0a96 | ||
|
|
38ecefaacb | ||
|
|
0e778296be | ||
|
|
6c9d8c7e41 | ||
|
|
1400ce105f | ||
|
|
154c865966 | ||
|
|
e8945c74de | ||
|
|
45c7252aed | ||
|
|
6146c64d81 | ||
|
|
ad7c8c21ea | ||
|
|
02a7b7fe48 | ||
|
|
2f145a98e0 | ||
|
|
5f4eeb054c | ||
|
|
680ce54dd4 | ||
|
|
fffce0a6b4 | ||
|
|
51b88b2265 | ||
|
|
b54cb272d0 | ||
|
|
d21e34e617 | ||
|
|
5a4b244e6b | ||
|
|
a6fd96f620 | ||
|
|
b03ceb806e | ||
|
|
25e0b725d1 | ||
|
|
1aba668a37 | ||
|
|
b53a266254 | ||
|
|
461e9becec | ||
|
|
9569fdfa36 | ||
|
|
8365c28cd5 | ||
|
|
4762a24022 | ||
|
|
57c7e0a8f8 | ||
|
|
393c6b236c | ||
|
|
4756971c88 | ||
|
|
5e794be8af | ||
|
|
73c8dae60d | ||
|
|
dc4dd898b7 | ||
|
|
bb1f376ae6 | ||
|
|
7e06d3ebba | ||
|
|
bb59eed82f | ||
|
|
cc038b31b6 | ||
|
|
a531a649fb | ||
|
|
8d703a6369 | ||
|
|
0dad6cc518 | ||
|
|
cff1065f5e | ||
|
|
ef05178855 | ||
|
|
87707ef0b8 | ||
|
|
825f148469 | ||
|
|
f82b16a0e9 | ||
|
|
7487c13b61 | ||
|
|
54c15d74a4 | ||
|
|
dbbc261075 | ||
|
|
f1108f1cbe | ||
|
|
812f485cd7 | ||
|
|
3c5b8bf50c | ||
|
|
525f80e0d2 | ||
|
|
edffc246ed | ||
|
|
7733c217c5 | ||
|
|
d917895569 | ||
|
|
158506b91e | ||
|
|
328bfe6b9b | ||
|
|
5b12764b83 | ||
|
|
53655a4ee5 | ||
|
|
6b808c5fe6 | ||
|
|
2a72b00679 | ||
|
|
c7b03457d7 | ||
|
|
494bb12500 | ||
|
|
419e997187 | ||
|
|
84d2d047ea | ||
|
|
122a50fe8c | ||
|
|
e555748807 | ||
|
|
f732f66709 | ||
|
|
82e037aad5 | ||
|
|
146c31586d | ||
|
|
df1c183e46 | ||
|
|
d01e3d7719 | ||
|
|
b63bd02969 | ||
|
|
57e8bf61e8 | ||
|
|
72e010d816 | ||
|
|
f1bd06134d | ||
|
|
ef0ef705fe | ||
|
|
d8855ec266 | ||
|
|
b8a74c1569 | ||
|
|
a388d2cb1a | ||
|
|
65397bfdeb | ||
|
|
ae0edc8a67 | ||
|
|
e1fef895b1 | ||
|
|
3a9db08b49 | ||
|
|
bdb3afd566 | ||
|
|
9fcc87761e |
@@ -28,7 +28,7 @@ jobs:
|
||||
# since sudo is required for usbgpu on macos, move the cache to a new location, as some of the files are owned by root
|
||||
PYTHONPYCACHEPREFIX: /tmp/tiny_python_pycache
|
||||
runs-on: [self-hosted, macOS]
|
||||
timeout-minutes: 20
|
||||
timeout-minutes: 60
|
||||
defaults:
|
||||
run:
|
||||
shell: bash -e -o pipefail {0}
|
||||
@@ -109,7 +109,7 @@ jobs:
|
||||
- name: Train MNIST
|
||||
run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
|
||||
- name: Run 10 CIFAR training steps
|
||||
run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=320 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
|
||||
run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=330 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
|
||||
- name: Run 10 CIFAR training steps w HALF
|
||||
run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=385 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
|
||||
#- name: Run 10 CIFAR training steps w BF16
|
||||
@@ -160,7 +160,7 @@ jobs:
|
||||
testnvidiabenchmark:
|
||||
name: tinybox green Benchmark
|
||||
runs-on: [self-hosted, Linux, tinyboxgreen]
|
||||
timeout-minutes: 30
|
||||
timeout-minutes: 60
|
||||
defaults:
|
||||
run:
|
||||
shell: bash -e -o pipefail {0}
|
||||
@@ -274,7 +274,7 @@ jobs:
|
||||
testmorenvidiabenchmark:
|
||||
name: tinybox green Training Benchmark
|
||||
runs-on: [self-hosted, Linux, tinyboxgreen]
|
||||
timeout-minutes: 20
|
||||
timeout-minutes: 60
|
||||
defaults:
|
||||
run:
|
||||
shell: bash -e -o pipefail {0}
|
||||
@@ -346,7 +346,7 @@ jobs:
|
||||
testamdbenchmark:
|
||||
name: tinybox red Benchmark
|
||||
runs-on: [self-hosted, Linux, tinybox]
|
||||
timeout-minutes: 20
|
||||
timeout-minutes: 60
|
||||
defaults:
|
||||
run:
|
||||
shell: bash -e -o pipefail {0}
|
||||
@@ -476,7 +476,7 @@ jobs:
|
||||
testmoreamdbenchmark:
|
||||
name: tinybox red Training Benchmark
|
||||
runs-on: [self-hosted, Linux, tinybox]
|
||||
timeout-minutes: 30
|
||||
timeout-minutes: 60
|
||||
defaults:
|
||||
run:
|
||||
shell: bash -e -o pipefail {0}
|
||||
@@ -511,8 +511,8 @@ jobs:
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=85 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
|
||||
- name: Run 10 CIFAR training steps w HALF
|
||||
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=188 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
|
||||
- name: Run 10 CIFAR training steps w BF16
|
||||
run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
|
||||
# - name: Run 10 CIFAR training steps w BF16
|
||||
# run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
|
||||
- name: Run 10 CIFAR training steps w winograd
|
||||
run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
|
||||
- name: Run full CIFAR training w 1 GPU
|
||||
@@ -539,7 +539,7 @@ jobs:
|
||||
testmlperfamdbenchmark:
|
||||
name: tinybox red MLPerf Benchmark
|
||||
runs-on: [self-hosted, Linux, tinybox]
|
||||
timeout-minutes: 30
|
||||
timeout-minutes: 60
|
||||
defaults:
|
||||
run:
|
||||
shell: bash -e -o pipefail {0}
|
||||
@@ -645,7 +645,7 @@ jobs:
|
||||
testreddriverbenchmark:
|
||||
name: AM Benchmark
|
||||
runs-on: [self-hosted, Linux, tinyboxrandom]
|
||||
timeout-minutes: 15
|
||||
timeout-minutes: 20
|
||||
defaults:
|
||||
run:
|
||||
shell: bash -e -o pipefail {0}
|
||||
@@ -716,7 +716,7 @@ jobs:
|
||||
testgreendriverbenchmark:
|
||||
name: NV Benchmark
|
||||
runs-on: [self-hosted, Linux, tinyboxrandom]
|
||||
timeout-minutes: 15
|
||||
timeout-minutes: 20
|
||||
defaults:
|
||||
run:
|
||||
shell: bash -e -o pipefail {0}
|
||||
|
||||
+108
-34
@@ -30,8 +30,6 @@ jobs:
|
||||
key: llvm-speed
|
||||
deps: testing_minimal
|
||||
llvm: 'true'
|
||||
- name: External Benchmark Schedule
|
||||
run: python3 test/external/external_benchmark_schedule.py
|
||||
- name: Speed Test
|
||||
run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
|
||||
- name: Speed Test (BEAM=2)
|
||||
@@ -48,7 +46,7 @@ jobs:
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
deps: docs
|
||||
pydeps: "capstone"
|
||||
pydeps: "capstone torch"
|
||||
- name: Build wheel and show size
|
||||
run: |
|
||||
pip install build
|
||||
@@ -79,6 +77,8 @@ jobs:
|
||||
run: |
|
||||
python docs/abstractions2.py
|
||||
python docs/abstractions3.py
|
||||
- name: Test README
|
||||
run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py && python README.py
|
||||
- name: Test Quickstart
|
||||
run: awk '/```python/{flag=1;next}/```/{flag=0}flag' docs/quickstart.md > quickstart.py && python quickstart.py
|
||||
- name: Test DEBUG
|
||||
@@ -160,8 +160,10 @@ jobs:
|
||||
with:
|
||||
key: be-minimal
|
||||
deps: testing_minimal
|
||||
- name: Test dtype with Python emulator
|
||||
run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
|
||||
- name: Test dtype with Python emulator (with RANGEIFY)
|
||||
run: |
|
||||
RANGEIFY=0 DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
|
||||
RANGEIFY=1 DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
|
||||
- name: Test ops with Python emulator
|
||||
run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20
|
||||
- name: Test uops with Python emulator
|
||||
@@ -259,21 +261,25 @@ jobs:
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: unittest-12
|
||||
pydeps: "pillow"
|
||||
pydeps: "pillow numpy ftfy regex"
|
||||
deps: testing_unit
|
||||
- name: Test README
|
||||
run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py && python README.py
|
||||
- name: Check Device.DEFAULT
|
||||
run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
|
||||
- name: Run unit tests
|
||||
run: python -m pytest -n=auto test/unit/ --durations=20
|
||||
run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
|
||||
- name: Run targetted tests on NULL backend
|
||||
run: NULL=1 python3 test/test_multitensor.py TestMultiTensor.test_data_parallel_resnet_train_step
|
||||
run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
|
||||
- name: Run SDXL on NULL backend
|
||||
run: MAX_BUFFER_SIZE=0 NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
|
||||
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
|
||||
# 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
|
||||
- name: Run GC tests
|
||||
run: python test/external/external_uop_gc.py
|
||||
- name: External Benchmark Schedule
|
||||
run: python3 test/external/external_benchmark_schedule.py
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
- name: Regen dataset on test_tiny
|
||||
@@ -310,9 +316,9 @@ jobs:
|
||||
run: python test/external/fuzz_shape_ops.py
|
||||
|
||||
testopenclimage:
|
||||
name: 'CL IMAGE Tests'
|
||||
name: CL IMAGE Tests
|
||||
runs-on: ubuntu-22.04
|
||||
timeout-minutes: 10
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
@@ -326,11 +332,15 @@ jobs:
|
||||
run: |
|
||||
CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
|
||||
CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
|
||||
- name: Test CL IMAGE=2 ops + training (rangeify)
|
||||
run: |
|
||||
RANGEIFY=1 CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
|
||||
RANGEIFY=1 CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
testgpumisc:
|
||||
name: 'CL Misc tests'
|
||||
name: CL Misc tests
|
||||
runs-on: ubuntu-22.04
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
@@ -355,7 +365,7 @@ jobs:
|
||||
path: /tmp/sops.gz
|
||||
|
||||
testopenpilot:
|
||||
name: 'openpilot Compile Tests'
|
||||
name: openpilot Compile Tests
|
||||
runs-on: ubuntu-22.04
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
@@ -370,7 +380,10 @@ jobs:
|
||||
llvm: 'true'
|
||||
- name: Test openpilot model kernel count and gate usage
|
||||
run: |
|
||||
ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2175 ALLOWED_GATED_READ_IMAGE=16 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=208 ALLOWED_READ_IMAGE=2160 ALLOWED_GATED_READ_IMAGE=16 RANGEIFY=0 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 model with rangeify
|
||||
run: |
|
||||
ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=33 RANGEIFY=1 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)
|
||||
@@ -387,7 +400,7 @@ jobs:
|
||||
# ****** ONNX Tests ******
|
||||
|
||||
testonnxcpu:
|
||||
name: 'ONNX (CPU) Tests'
|
||||
name: ONNX (CPU) Tests
|
||||
runs-on: ubuntu-22.04
|
||||
timeout-minutes: 20
|
||||
|
||||
@@ -415,7 +428,7 @@ jobs:
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
testopencl:
|
||||
name: 'ONNX (GPU)+Optimization Tests'
|
||||
name: ONNX (CL)+Optimization Tests
|
||||
runs-on: ubuntu-22.04
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
@@ -439,8 +452,12 @@ jobs:
|
||||
run: CL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
|
||||
- name: Test MLPerf stuff
|
||||
run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
|
||||
- name: NULL=1 beautiful_mnist_multigpu
|
||||
run: NULL=1 python examples/beautiful_mnist_multigpu.py
|
||||
- name: Test Bert training
|
||||
run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
|
||||
- name: Test llama 3 training
|
||||
run: MAX_BUFFER_SIZE=0 DEV=NULL SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
|
||||
run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
@@ -503,8 +520,8 @@ jobs:
|
||||
|
||||
# ****** Feature Tests ******
|
||||
|
||||
testrangeify:
|
||||
name: Linux (rangeify)
|
||||
testrangeifycpu:
|
||||
name: Linux (rangeify) CPU
|
||||
runs-on: ubuntu-24.04
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
@@ -519,22 +536,71 @@ jobs:
|
||||
llvm: "true"
|
||||
- name: Test CPU=1 RANGEIFY=1
|
||||
# TODO: add more passing tests here
|
||||
# test_symbolic_arange_sym_step is passing now
|
||||
# test_threefry_doesnt_use_long is because there's a contig after the long now
|
||||
run: |
|
||||
CPU=1 CPU_LLVM=0 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \
|
||||
-k "not test_symbolic_arange_sym_step and not test_threefry_doesnt_use_long" \
|
||||
test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_tensor_variable.py \
|
||||
test/test_outerworld_range.py test/test_sample.py test/test_randomness.py
|
||||
- name: Test multitensor
|
||||
run: RANGEIFY=1 PYTHONPATH="." python3 test/test_multitensor.py TestMultiTensor.test_matmul_shard_1_1 TestMultiTensor.test_simple_add_W
|
||||
- name: Test GPU=1 RANGEIFY=1
|
||||
run: GPU=1 RANGEIFY=1 pytest -n auto test/test_ops.py
|
||||
- name: Test CPU=1 RANGEIFY=2
|
||||
run: CPU=1 CPU_LLVM=0 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
|
||||
test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_symbolic_ops.py test/test_symbolic_jit.py test/test_tensor_variable.py \
|
||||
test/test_outerworld_range.py test/test_randomness.py test/test_nn.py test/test_arange.py test/test_tensor.py test/test_optim.py \
|
||||
test/test_setitem.py test/test_assign.py test/test_multitensor.py test/test_const_folding.py
|
||||
- name: Test CPU=1 DEVECTORIZE=0 (RANGEIFY=1)
|
||||
run: CPU=1 CPU_LLVM=0 RANGEIFY=1 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
|
||||
- name: Test CPU=1 CPU_LLVM=1 RANGEIFY=1
|
||||
run: |
|
||||
CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_edgecases.py
|
||||
- name: Test Docs RANGEIFY=1
|
||||
run: |
|
||||
RANGEIFY=1 python docs/abstractions2.py
|
||||
# RANGEIFY=2 isn't supported
|
||||
#- name: Test CPU=1 RANGEIFY=2
|
||||
# run: CPU=1 CPU_LLVM=0 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
|
||||
# slow (and still wrong on beautiful_mnist)
|
||||
#- name: Test LLVM=1 RANGEIFY=1 (slow tests)
|
||||
#- name: Test LLVM RANGEIFY=1 (slow tests)
|
||||
# run: CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
testrangeifycl:
|
||||
name: Linux (rangeify) CL
|
||||
runs-on: ubuntu-24.04
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: rangeify-cl
|
||||
deps: testing
|
||||
opencl: 'true'
|
||||
llvm: "true"
|
||||
- name: Test CL=1 RANGEIFY=1
|
||||
run: CL=1 RANGEIFY=1 pytest -n auto test/test_ops.py test/test_schedule.py test/test_symbolic_ops.py test/test_jit.py test/unit/test_disk_tensor.py test/models/test_mnist.py test/unit/test_mnist_dataset.py test/test_optim.py --durations 20
|
||||
- name: Test Fuse
|
||||
run: CL=1 RANGEIFY=2 python3 -m pytest --durations 20 test/test_softmax_fusion.py -k "not test_auto_softmax"
|
||||
- name: Test ONNX
|
||||
run: CL=1 RANGEIFY=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
testrangeifymacos:
|
||||
name: MacOS (rangeify)
|
||||
runs-on: macos-14
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: metal
|
||||
deps: testing
|
||||
- name: some unit tests
|
||||
run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/unit/test_winograd.py test/unit/test_linalg.py --durations=20
|
||||
- name: Test METAL=1 RANGEIFY=1
|
||||
run: |
|
||||
METAL=1 RANGEIFY=1 python -m pytest -n=auto test/test_ops.py test/test_multitensor.py --durations=20
|
||||
METAL=1 MAX_KERNEL_BUFFERS=6 RANGEIFY=1 PYTHONPATH=. python test/test_multitensor.py TestBatchNorm.test_batchnorm
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
testdevectorize:
|
||||
name: Linux (devectorize)
|
||||
@@ -656,8 +722,10 @@ jobs:
|
||||
run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
|
||||
- name: Run TestOps.test_add with SQTT
|
||||
run: |
|
||||
PROFILE=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
|
||||
VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
|
||||
extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
|
||||
- name: Run pytest (amd) with RANGEIFY
|
||||
run: RANGEIFY=1 python -m pytest test/test_linearizer.py::TestLinearizer::test_where_fold
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
@@ -977,3 +1045,9 @@ jobs:
|
||||
run: |
|
||||
python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
|
||||
python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
|
||||
- name: Run pytest (${{ matrix.backend }}) with RANGEIFY
|
||||
if: matrix.backend=='webgpu'
|
||||
env:
|
||||
RANGEIFY: 1
|
||||
shell: bash
|
||||
run: python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
|
||||
|
||||
@@ -80,7 +80,9 @@ print("******** third, the UOp ***********")
|
||||
|
||||
from tinygrad.engine.realize import run_schedule
|
||||
from tinygrad.engine.schedule import create_schedule_with_vars
|
||||
from tinygrad.helpers import RANGEIFY
|
||||
from tinygrad.schedule.kernelize import get_kernelize_map
|
||||
from tinygrad.schedule.rangeify import get_rangeify_map
|
||||
|
||||
# allocate some values + load in values
|
||||
a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
|
||||
@@ -93,10 +95,10 @@ out = a + b
|
||||
s = UOp(Ops.SINK, dtypes.void, (out,))
|
||||
|
||||
# group the computation into kernels
|
||||
becomes_map = get_kernelize_map(s)
|
||||
becomes_map = get_rangeify_map(s) if RANGEIFY else get_kernelize_map(s)
|
||||
|
||||
# the compute maps to an assign
|
||||
assign = becomes_map[a+b]
|
||||
assign = becomes_map[a+b].base
|
||||
|
||||
# the first source is the output buffer (data)
|
||||
assert assign.src[0].op is Ops.BUFFER
|
||||
|
||||
@@ -41,8 +41,6 @@ BEAM | [#] | number of beams in kernel beam search
|
||||
DEFAULT_FLOAT | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
|
||||
IMAGE | [1-2] | enable 2d specific optimizations
|
||||
FLOAT16 | [1] | use float16 for images instead of float32
|
||||
PTX | [1] | enable the specialized [PTX](https://docs.nvidia.com/cuda/parallel-thread-execution/) assembler for Nvidia GPUs. If not set, defaults to generic CUDA codegen backend.
|
||||
PROFILE | [1] | enable profiling. This feature is supported in NV, AMD, QCOM and METAL backends.
|
||||
VISIBLE_DEVICES | [list[int]]| restricts the NV/AMD devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
|
||||
JIT | [0-2] | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
|
||||
VIZ | [1] | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
|
||||
|
||||
+18
-11
@@ -2,17 +2,17 @@
|
||||
|
||||
tinygrad supports various runtimes, enabling your code to scale across a wide range of devices. The default runtime can be automatically selected based on the available hardware, or you can force a specific runtime to be default using environment variables (e.g., `CPU=1`).
|
||||
|
||||
| Runtime | Description | Requirements |
|
||||
|---------|-------------|--------------|
|
||||
| [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | Ampere/Ada series GPUs |
|
||||
| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | RDNA2/RDNA3/RDNA4 series GPUs. You can select one of the interfaces for communication by setting `AMD_IFACE=(KFD|PCI)`. See [AMD interfaces](#amd-interfaces) for more details. |
|
||||
| [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | 6xx series GPUs |
|
||||
| [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | M1+ Macs; Metal 3.0+ for `bfloat` support |
|
||||
| [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | NVIDIA GPU with CUDA support |
|
||||
| [OpenCL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | OpenCL 2.0 compatible device |
|
||||
| [CPU (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` |
|
||||
| [LLVM (LLVM IR)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | llvm libraries installed and findable |
|
||||
| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0). |
|
||||
| Runtime | Description | Compiler Options | Requirements |
|
||||
|---------|-------------|------------------|--------------|
|
||||
| [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | nvrtc (default)<br>PTX (`NV_PTX=1`) | Ampere/Ada/Blackwell series GPUs.<br>You can select an interface via `NV_IFACE=(NVK\|PCI)`. See [NV interfaces](#nv-interfaces) for details. |
|
||||
| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`AMD_LLVM=1`)<br>HIP/COMGR (`AMD_HIP=1`) | RDNA2 or newer GPUs.<br>You can select an interface via `AMD_IFACE=(KFD\|PCI\|USB)`. See [AMD interfaces](#amd-interfaces) for details. |
|
||||
| [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | - | 6xx series GPUs |
|
||||
| [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | - | M1+ Macs; Metal 3.0+ for `bfloat` support |
|
||||
| [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | nvrtc (default)<br> PTX (`CUDA_PTX=1`) | NVIDIA GPU with CUDA support |
|
||||
| [CL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | - | OpenCL 2.0 compatible device |
|
||||
| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`CPU_LLVM=1`) | `clang` compiler in system `PATH` |
|
||||
| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | - | Dawn library installed and discoverable. Binaries: [pydawn v0.3.0](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0) |
|
||||
|
||||
|
||||
## Interoperability
|
||||
|
||||
@@ -70,5 +70,12 @@ AMD backend supports several interfaces for communicating with devices:
|
||||
|
||||
* `KFD`: uses the amdgpu driver
|
||||
* `PCI`: uses the [AM driver](developer/am.md)
|
||||
* `USB`: USB3 interafce for asm24xx chips.
|
||||
|
||||
You can force an interface by setting `AMD_IFACE` to one of these values. In the case of `AMD_IFACE=PCI`, this may unbind your GPU from the amdgpu driver.
|
||||
|
||||
## NV Interfaces
|
||||
NV backend supports several interfaces for communicating with devices:
|
||||
|
||||
* `NVK`: uses the nvidia driver
|
||||
* `PCI`: uses the [NV driver](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/support/nv/nvdev.py)
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import sys, time
|
||||
from tinygrad import TinyJit, GlobalCounters, fetch, getenv
|
||||
from tinygrad.frontend.onnx import OnnxRunner
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from extra.onnx_helpers import get_example_inputs, validate
|
||||
|
||||
def load_onnx_model(onnx_file):
|
||||
|
||||
@@ -8,7 +8,7 @@ import numpy as np
|
||||
import subprocess
|
||||
import tensorflow as tf
|
||||
import tf2onnx
|
||||
from tinygrad.frontend.onnx import OnnxRunner
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.helpers import to_mv
|
||||
from extra.export_model import export_model_clang, compile_net, jit_model
|
||||
|
||||
+6
-6
@@ -26,8 +26,8 @@ class Attention:
|
||||
start_pos = start_pos.val
|
||||
|
||||
if HALF: x = x.half()
|
||||
xqkv = self.c_attn(x)
|
||||
xq, xk, xv = [xqkv.shrink((None, None, (i*self.dim, (i+1)*self.dim))).reshape(None, None, self.n_heads, self.head_dim) for i in range(3)]
|
||||
xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
|
||||
xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
|
||||
bsz, seqlen, _, _ = xq.shape
|
||||
|
||||
# create kv cache
|
||||
@@ -35,11 +35,11 @@ class Attention:
|
||||
self.cache_kv = Tensor.zeros(2, bsz, MAX_CONTEXT, self.n_heads, self.head_dim, dtype=x.dtype).contiguous().realize()
|
||||
|
||||
# update the cache
|
||||
self.cache_kv.shrink((None, None,(start_pos,start_pos+seqlen),None,None)).assign(Tensor.stack(xk, xv)).realize()
|
||||
self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()
|
||||
|
||||
if start_pos > 0:
|
||||
keys = self.cache_kv[0].shrink((None, (0, start_pos+seqlen), None, None))
|
||||
values = self.cache_kv[1].shrink((None, (0, start_pos+seqlen), None, None))
|
||||
keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
|
||||
values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
|
||||
else:
|
||||
keys = xk
|
||||
values = xv
|
||||
@@ -64,7 +64,7 @@ class TransformerBlock:
|
||||
|
||||
def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]):
|
||||
h = x + self.attn(self.ln_1(x), start_pos, mask).float()
|
||||
return (h + self.mlp(self.ln_2(h)))
|
||||
return (h + self.mlp(self.ln_2(h))).contiguous()
|
||||
|
||||
class Transformer:
|
||||
def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):
|
||||
|
||||
@@ -229,7 +229,8 @@ def train_cifar():
|
||||
if getenv("RANDOM_CROP", 1):
|
||||
X = random_crop(X, crop_size=32)
|
||||
if getenv("RANDOM_FLIP", 1):
|
||||
X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X) # flip LR
|
||||
# NOTE: RANGEIFY=1 needs this contiguous or the X[perms] is very slow
|
||||
X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X).contiguous() # flip LR
|
||||
X, Y = X[perms], Y[perms]
|
||||
return X, Y, *cutmix(X, Y, perms, mask_size=hyp['net']['cutmix_size'])
|
||||
|
||||
|
||||
@@ -511,6 +511,33 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
|
||||
# happens with BENCHMARK set
|
||||
pass
|
||||
|
||||
# stable diffusion callbacks to match mlperf ref; declared here because they're pickled
|
||||
def filter_dataset(sample:dict): return {k:v for k,v in sample.items() if k in {'npy', 'txt'}}
|
||||
def collate(batch:list[dict]):
|
||||
ret = {"npy": [], "txt": [], "__key__": []}
|
||||
for sample in batch:
|
||||
for k,v in sample.items():
|
||||
ret[k].append(v)
|
||||
return ret
|
||||
def collate_fn(batch): return batch
|
||||
|
||||
# Reference (code): https://github.com/mlcommons/training/blob/2f4a93fb4888180755a8ef55f4b977ef8f60a89e/stable_diffusion/ldm/data/webdatasets.py, Line 55
|
||||
# Reference (params): https://github.com/mlcommons/training/blob/ab4ae1ca718d7fe62c369710a316dff18768d04b/stable_diffusion/configs/train_01x08x08.yaml, Line 107
|
||||
def batch_load_train_stable_diffusion(urls:str, BS:int):
|
||||
import webdataset
|
||||
dataset = webdataset.WebDataset(urls=urls, resampled=True, cache_size=-1, cache_dir=None)
|
||||
dataset = dataset.shuffle(size=1000)
|
||||
dataset = dataset.decode()
|
||||
dataset = dataset.map(filter_dataset)
|
||||
dataset = dataset.batched(BS, partial=False, collation_fn=collate)
|
||||
dataset = webdataset.WebLoader(dataset, batch_size=None, shuffle=False, num_workers=1, persistent_workers=True, collate_fn=collate_fn)
|
||||
|
||||
for x in dataset:
|
||||
assert isinstance(x, dict) and all(isinstance(k, str) for k in x.keys()) and all(isinstance(v, list) for v in x.values())
|
||||
assert all(isinstance(moment_mean_logvar, np.ndarray) and moment_mean_logvar.shape==(1,8,64,64) for moment_mean_logvar in x["npy"])
|
||||
assert all(isinstance(caption, str) for caption in x["txt"])
|
||||
yield x
|
||||
|
||||
# llama3
|
||||
|
||||
class BinIdxDataset:
|
||||
|
||||
@@ -2,7 +2,9 @@ import math
|
||||
from typing import Union
|
||||
|
||||
from tinygrad import Tensor, nn, dtypes
|
||||
from tinygrad.helpers import prod, argfix
|
||||
from tinygrad.helpers import prod, argfix, Context
|
||||
from tinygrad.nn.state import get_parameters
|
||||
from extra.models.unet import UNetModel
|
||||
|
||||
# rejection sampling truncated randn
|
||||
def rand_truncn(*shape, dtype=None, truncstds=2, **kwargs) -> Tensor:
|
||||
@@ -17,6 +19,10 @@ def he_normal(*shape, a: float = 0.00, **kwargs) -> Tensor:
|
||||
std = math.sqrt(2.0 / (1 + a ** 2)) / math.sqrt(prod(argfix(*shape)[1:])) / 0.87962566103423978
|
||||
return std * rand_truncn(*shape, **kwargs)
|
||||
|
||||
# Stable Diffusion v2 training uses default torch gelu, which doesn't use tanh approximation
|
||||
def gelu_erf(x:Tensor) -> Tensor:
|
||||
return 0.5 * x * (1.0 + (x / 1.4142135623730951).erf())
|
||||
|
||||
class Conv2dHeNormal(nn.Conv2d):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
|
||||
super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
|
||||
@@ -127,3 +133,59 @@ class Conv2dRetinaNet(nn.Conv2d):
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
|
||||
groups=self.groups, stride=self.stride, dilation=self.dilation, padding=self.padding)
|
||||
|
||||
# copy torch AMP: isolate mixed precision to just the below autocast ops, instead of using dtypes.default_float which affects all new Tensors
|
||||
class AutocastLinear(nn.Linear):
|
||||
cast_dtype=dtypes.bfloat16 # enable monkeypatching of the mixed precision dtype
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
dtype = type(self).cast_dtype
|
||||
return x.cast(dtype).linear(self.weight.cast(dtype).transpose(), self.bias.cast(dtype) if self.bias is not None else None)
|
||||
|
||||
class AutocastConv2d(nn.Conv2d):
|
||||
cast_dtype=dtypes.bfloat16
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
dtype = type(self).cast_dtype
|
||||
return x.cast(dtype).conv2d(self.weight.cast(dtype), self.bias.cast(dtype), self.groups, self.stride, self.dilation, self.padding)
|
||||
|
||||
# copy torch AMP: upcast to float32 before GroupNorm and LayerNorm
|
||||
class AutocastGroupNorm(nn.GroupNorm):
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
return super().__call__(x.cast(dtypes.float32))
|
||||
|
||||
class AutocastLayerNorm(nn.LayerNorm):
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
return super().__call__(x.cast(dtypes.float32))
|
||||
|
||||
def zero_module(module):
|
||||
for p in get_parameters(module): p.assign(Tensor.zeros_like(p).contiguous())
|
||||
|
||||
# Stable Diffusion mlperf reference doesn't call scaled_dot_product_attention
|
||||
# copy torch AMP: upcast to float32 before softmax on CUDA
|
||||
def attn_f32_softmax(q:Tensor, k:Tensor, v:Tensor) -> Tensor:
|
||||
return (q.matmul(k.transpose(-2,-1), dtype=dtypes.float32) / math.sqrt(q.shape[-1])).softmax(-1).cast(q.dtype) @ v
|
||||
|
||||
def init_stable_diffusion(version:str, pretrained:str, devices:list[str]):
|
||||
from examples.stable_diffusion import StableDiffusion
|
||||
from tinygrad.nn.state import safe_load, safe_save, load_state_dict, get_state_dict
|
||||
from tempfile import TemporaryDirectory
|
||||
model = StableDiffusion(version=version, pretrained=pretrained)
|
||||
unet:UNetModel = model.model.diffusion_model
|
||||
|
||||
# this prevents extra consumption of memory, enabling much larger BS
|
||||
Tensor.realize(*get_parameters(unet))
|
||||
with TemporaryDirectory(prefix="unet_init") as tmp:
|
||||
safe_save(get_state_dict(unet), init_fn:=f"{tmp}/init_model.safetensors")
|
||||
load_state_dict(unet, safe_load(init_fn))
|
||||
|
||||
sqrt_alphas_cumprod = model.alphas_cumprod.sqrt().realize()
|
||||
sqrt_one_minus_alphas_cumprod = (1 - model.alphas_cumprod).sqrt().realize()
|
||||
|
||||
if len(devices) > 1:
|
||||
to_move = [sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod]
|
||||
if version == "v2-mlperf-train": to_move += get_parameters(unet) + get_parameters(model.cond_stage_model)
|
||||
for p in to_move:
|
||||
p.to_(devices)
|
||||
with Context(BEAM=0):
|
||||
Tensor.realize(*to_move)
|
||||
|
||||
return model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
import math
|
||||
from tinygrad import dtypes
|
||||
from tinygrad import dtypes, Tensor
|
||||
from tinygrad.nn.optim import Optimizer
|
||||
|
||||
from extra.lr_scheduler import LR_Scheduler
|
||||
from typing import Callable
|
||||
|
||||
# https://github.com/mlcommons/training/blob/e237206991d10449d9675d95606459a3cb6c21ad/image_classification/tensorflow2/lars_util.py
|
||||
class PolynomialDecayWithWarmup(LR_Scheduler):
|
||||
@@ -36,4 +37,24 @@ class CosineAnnealingLRWithWarmup(LR_Scheduler):
|
||||
def get_lr(self):
|
||||
warmup_lr = ((self.epoch_counter+1) / self.warmup_steps) * self.base_lr
|
||||
decay_lr = self.end_lr + 0.5 * (self.base_lr-self.end_lr) * (1 + (((self.epoch_counter+1-self.warmup_steps)/self.decay_steps) * math.pi).cos())
|
||||
return (self.epoch_counter < self.warmup_steps).where(warmup_lr, decay_lr).cast(self.optimizer.lr.dtype)
|
||||
return (self.epoch_counter < self.warmup_steps).where(warmup_lr, decay_lr).cast(self.optimizer.lr.dtype)
|
||||
|
||||
# Reference: https://github.com/mlcommons/training/blob/64b14a9abc74e08779a175abca7d291f8c957632/stable_diffusion/ldm/lr_scheduler.py, Lines 36-97
|
||||
class LambdaLinearScheduler:
|
||||
def __init__(self, warm_up_steps:int, f_min:float, f_max:float, f_start:float, cycle_lengths:int):
|
||||
self.lr_warm_up_steps, self.f_min, self.f_max, self.f_start, self.cycle_lengths = warm_up_steps, f_min, f_max, f_start, cycle_lengths
|
||||
|
||||
def schedule(self, n:Tensor) -> Tensor:
|
||||
warm_up = (n < self.lr_warm_up_steps)
|
||||
f_warm_up = (self.f_max - self.f_start) / self.lr_warm_up_steps * n + self.f_start
|
||||
return warm_up.where(f_warm_up, self.f_min + (self.f_max - self.f_min) * (self.cycle_lengths - n) / (self.cycle_lengths))
|
||||
|
||||
# based on torch.optim.lr_scheduler.LambdaLR
|
||||
class LambdaLR(LR_Scheduler):
|
||||
def __init__(self, optimizer:Optimizer, base_lr:Tensor, lr_lambda:Callable):
|
||||
super().__init__(optimizer)
|
||||
self.base_lr, self.lr_lambda = base_lr, lr_lambda
|
||||
self.step()
|
||||
|
||||
def get_lr(self):
|
||||
return self.base_lr * self.lr_lambda(self.epoch_counter - 1)
|
||||
@@ -1,10 +1,10 @@
|
||||
import time, math
|
||||
import time, math, os
|
||||
start = time.perf_counter()
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, Device, dtypes, GlobalCounters, TinyJit
|
||||
from tinygrad.nn.state import get_parameters, load_state_dict, safe_load
|
||||
from tinygrad.helpers import getenv
|
||||
from tinygrad.helpers import getenv, Context, prod
|
||||
from extra.bench_log import BenchEvent, WallTimeEvent
|
||||
def tlog(x): print(f"{x:25s} @ {time.perf_counter()-start:5.2f}s")
|
||||
|
||||
@@ -287,6 +287,256 @@ def eval_llama3():
|
||||
log_perplexity = np.mean(losses)
|
||||
print(f"Log Perplexity: {log_perplexity}")
|
||||
|
||||
# NOTE: BEAM hangs on 8xmi300x with DECODE_BS=384 in final realize below; function is declared here for external testing
|
||||
@TinyJit
|
||||
def vae_decode(x:Tensor, vae, disable_beam=False) -> Tensor:
|
||||
from examples.stable_diffusion import AutoencoderKL
|
||||
assert isinstance(vae, AutoencoderKL)
|
||||
x = vae.post_quant_conv(1./0.18215 * x)
|
||||
|
||||
x = vae.decoder.conv_in(x)
|
||||
x = vae.decoder.mid(x)
|
||||
for i, l in enumerate(vae.decoder.up[::-1]):
|
||||
print("decode", x.shape)
|
||||
for b in l['block']: x = b(x)
|
||||
if 'upsample' in l:
|
||||
bs,c,py,px = x.shape
|
||||
x = x.reshape(bs, c, py, 1, px, 1).expand(bs, c, py, 2, px, 2).reshape(bs, c, py*2, px*2)
|
||||
x = l['upsample']['conv'](x)
|
||||
if i == len(vae.decoder.up) - 1 and disable_beam:
|
||||
with Context(BEAM=0): x.realize()
|
||||
else: x.realize()
|
||||
x = vae.decoder.conv_out(vae.decoder.norm_out(x).swish())
|
||||
|
||||
x = ((x + 1.0) / 2.0).clip(0.0, 1.0)
|
||||
return x
|
||||
|
||||
def eval_stable_diffusion():
|
||||
import csv, PIL, sys
|
||||
from tqdm import tqdm
|
||||
from examples.mlperf.initializers import init_stable_diffusion, gelu_erf
|
||||
from examples.stable_diffusion import AutoencoderKL
|
||||
from extra.models.unet import UNetModel
|
||||
from tinygrad.nn.state import load_state_dict, torch_load
|
||||
from tinygrad.helpers import BEAM
|
||||
from extra.models import clip
|
||||
from extra.models.clip import FrozenOpenClipEmbedder
|
||||
from extra.models.clip import OpenClipEncoder
|
||||
from extra.models.inception import FidInceptionV3
|
||||
|
||||
config = {}
|
||||
GPUS = config["GPUS"] = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 1))]
|
||||
for x in GPUS: Device[x]
|
||||
print(f"running eval on {GPUS}")
|
||||
seed = config["seed"] = getenv("SEED", 12345)
|
||||
CKPTDIR = config["CKPTDIR"] = Path(getenv("CKPTDIR", "./checkpoints"))
|
||||
DATADIR = config["DATADIR"] = Path(getenv("DATADIR", "./datasets"))
|
||||
CONTEXT_BS = config["CONTEXT_BS"] = getenv("CONTEXT_BS", 1 * len(GPUS))
|
||||
DENOISE_BS = config["DENOISE_BS"] = getenv("DENOISE_BS", 1 * len(GPUS))
|
||||
DECODE_BS = config["DECODE_BS"] = getenv("DECODE_BS", 1 * len(GPUS))
|
||||
INCEPTION_BS = config["INCEPTION_BS"] = getenv("INCEPTION_BS", 1 * len(GPUS))
|
||||
CLIP_BS = config["CLIP_BS"] = getenv("CLIP_BS", 1 * len(GPUS))
|
||||
EVAL_CKPT_DIR = config["EVAL_CKPT_DIR"] = getenv("EVAL_CKPT_DIR", "")
|
||||
STOP_IF_CONVERGED = config["STOP_IF_CONVERGED"] = getenv("STOP_IF_CONVERGED", 0)
|
||||
|
||||
if (WANDB := getenv("WANDB", "")):
|
||||
import wandb
|
||||
wandb.init(config=config, project="MLPerf-Stable-Diffusion")
|
||||
|
||||
assert EVAL_CKPT_DIR != "", "provide a directory with checkpoints to be evaluated"
|
||||
print(f"running eval on checkpoints in {EVAL_CKPT_DIR}\nSEED={seed}")
|
||||
eval_queue:list[tuple[int, Path]] = []
|
||||
for p in Path(EVAL_CKPT_DIR).iterdir():
|
||||
if p.name.endswith(".safetensors"):
|
||||
ckpt_iteration = p.name.split(".safetensors")[0]
|
||||
assert ckpt_iteration.isdigit(), f"invalid checkpoint name: {p.name}, expected <digits>.safetensors"
|
||||
eval_queue.append((int(ckpt_iteration), p))
|
||||
assert len(eval_queue), f'no files ending with ".safetensors" were found in {EVAL_CKPT_DIR}'
|
||||
print(sorted(eval_queue, reverse=True))
|
||||
|
||||
Tensor.manual_seed(seed) # seed for weight initialization
|
||||
model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod = init_stable_diffusion("v2-mlperf-eval", CKPTDIR / "sd" / "512-base-ema.ckpt", GPUS)
|
||||
|
||||
# load prompts for generating images for validation; 2 MB of data total
|
||||
with open(DATADIR / "coco2014" / "val2014_30k.tsv") as f:
|
||||
reader = csv.DictReader(f, delimiter="\t")
|
||||
eval_inputs:list[dict] = [{"image_id": int(row["image_id"]), "id": int(row["id"]), "caption": row["caption"]} for row in reader]
|
||||
assert len(eval_inputs) == 30_000
|
||||
# NOTE: the clip weights are the same between model.cond_stage_model and clip_encoder
|
||||
eval_timesteps = list(reversed(range(1, 1000, 20)))
|
||||
|
||||
original_device, Device.DEFAULT = Device.DEFAULT, "CPU"
|
||||
# The choice of alphas_prev[0] = alphas_cumprod[0] seems arbitrary, but it's how the mlperf ref does it:
|
||||
# alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
|
||||
eval_alphas_prev = model.alphas_cumprod[0:1].cat(model.alphas_cumprod[list(range(1, 1000, 20))[:-1]]).to(GPUS).realize()
|
||||
inception = FidInceptionV3().load_from_pretrained(CKPTDIR / "inception" / "pt_inception-2015-12-05-6726825d.pth")
|
||||
vision_cfg = {'width': 1280, 'layers': 32, 'd_head': 80, 'image_size': 224, 'patch_size': 14}
|
||||
text_cfg = {'width': 1024, 'n_heads': 16, 'layers': 24, 'vocab_size': 49408, 'ctx_length': 77}
|
||||
clip.gelu = gelu_erf
|
||||
clip_encoder = OpenClipEncoder(1024, text_cfg, vision_cfg)
|
||||
loaded = torch_load(CKPTDIR / "clip" / "open_clip_pytorch_model.bin")
|
||||
loaded.update({"attn_mask": clip_encoder.attn_mask, "mean": clip_encoder.mean, "std": clip_encoder.std})
|
||||
load_state_dict(clip_encoder, loaded)
|
||||
Device.DEFAULT=original_device
|
||||
|
||||
@TinyJit
|
||||
def denoise_step(x:Tensor, x_x:Tensor, t_t:Tensor, uc_c:Tensor, sqrt_alphas_cumprod_t:Tensor, sqrt_one_minus_alphas_cumprod_t:Tensor,
|
||||
alpha_prev:Tensor, unet:UNetModel, GPUS) -> Tensor:
|
||||
out_uncond, out = unet(x_x, t_t, uc_c).to("CPU").reshape(-1, 2, 4, 64, 64).chunk(2, dim=1)
|
||||
out_uncond = out_uncond.squeeze(1).shard(GPUS,axis=0)
|
||||
out = out.squeeze(1).shard(GPUS,axis=0)
|
||||
v_t = out_uncond + 8.0 * (out - out_uncond)
|
||||
e_t = sqrt_alphas_cumprod_t * v_t + sqrt_one_minus_alphas_cumprod_t * x
|
||||
pred_x0 = sqrt_alphas_cumprod_t * x - sqrt_one_minus_alphas_cumprod_t * v_t
|
||||
dir_xt = (1. - alpha_prev).sqrt() * e_t
|
||||
x_prev = alpha_prev.sqrt() * pred_x0 + dir_xt
|
||||
return x_prev.realize()
|
||||
|
||||
def shard_tensor(t:Tensor) -> Tensor: return t.shard(GPUS, axis=0) if len(GPUS) > 1 else t.to(GPUS[0])
|
||||
def get_batch(whole:Tensor, i:int, bs:int) -> tuple[Tensor, int]:
|
||||
batch = whole[i: i + bs].to("CPU")
|
||||
if (unpadded_bs:=batch.shape[0]) < bs:
|
||||
batch = batch.cat(batch[-1:].expand(bs - unpadded_bs, *batch[-1].shape))
|
||||
return batch, unpadded_bs
|
||||
|
||||
@Tensor.train(mode=False)
|
||||
def eval_unet(eval_inputs:list[dict], unet:UNetModel, cond_stage:FrozenOpenClipEmbedder, first_stage:AutoencoderKL,
|
||||
inception:FidInceptionV3, clip:OpenClipEncoder) -> tuple[float, float]:
|
||||
# Eval is divided into 5 jits, one per model
|
||||
# It doesn't make sense to merge these jits, e.g. unet repeats 50 times in isolation; images fork to separate inception/clip
|
||||
# We're generating and scoring 30,000 images per eval, and all the data can flow through one jit at a time
|
||||
# To maximize throughput for each jit, we have only one model/jit on the GPU at a time, and pool outputs from each jit off-GPU
|
||||
for model in (unet, first_stage, inception, clip):
|
||||
Tensor.realize(*[p.to_("CPU") for p in get_parameters(model)])
|
||||
|
||||
uc_written = False
|
||||
models = (cond_stage, unet, first_stage, inception, clip)
|
||||
jits = (jit_context:=TinyJit(cond_stage.embed_tokens), denoise_step, vae_decode, jit_inception:=TinyJit(inception),
|
||||
jit_clip:=TinyJit(clip.get_clip_score))
|
||||
all_bs = (CONTEXT_BS, DENOISE_BS, DECODE_BS, INCEPTION_BS, CLIP_BS)
|
||||
if (EVAL_SAMPLES:=getenv("EVAL_SAMPLES", 0)) and EVAL_SAMPLES > 0:
|
||||
eval_inputs = eval_inputs[0:EVAL_SAMPLES]
|
||||
output_shapes = [(ns:=len(eval_inputs),77), (ns,77,1024), (ns,4,64,64), (ns,3,512,512), (ns,2048), (ns,)]
|
||||
# Writing progress to disk lets us resume eval if we crash
|
||||
stages = ["tokens", "embeds", "latents", "imgs", "inception", "clip"]
|
||||
disk_tensor_names, disk_tensor_shapes = stages + ["end", "uc"], output_shapes + [(6,), (1,77,1024)]
|
||||
if not all(os.path.exists(f"{EVAL_CKPT_DIR}/{name}.bytes") for name in disk_tensor_names):
|
||||
for name, shape in zip(disk_tensor_names, disk_tensor_shapes):
|
||||
file = Path(f"{EVAL_CKPT_DIR}/{name}.bytes")
|
||||
file.unlink(missing_ok=True)
|
||||
with file.open("wb") as f: f.truncate(prod(shape) * 4)
|
||||
progress = {name: Tensor.empty(*shape, device=f"disk:{EVAL_CKPT_DIR}/{name}.bytes", dtype=dtypes.int if name in {"tokens", "end"} else dtypes.float)
|
||||
for name, shape in zip(disk_tensor_names, disk_tensor_shapes)}
|
||||
|
||||
def embed_tokens(tokens:Tensor) -> Tensor:
|
||||
nonlocal uc_written
|
||||
if not uc_written:
|
||||
with Context(BEAM=0): progress["uc"].assign(cond_stage.embed_tokens(cond_stage.tokenize("").to(GPUS)).to("CPU").realize()).realize()
|
||||
uc_written = True
|
||||
return jit_context(shard_tensor(tokens))
|
||||
|
||||
def generate_latents(embeds:Tensor) -> Tensor:
|
||||
uc_c = Tensor.stack(progress["uc"].to("CPU").expand(bs, 77, 1024), embeds, dim=1).reshape(-1, 77, 1024)
|
||||
uc_c = shard_tensor(uc_c)
|
||||
x = shard_tensor(Tensor.randn(bs,4,64,64))
|
||||
for step_idx, timestep in enumerate(tqdm(eval_timesteps)):
|
||||
reversed_idx = Tensor([50 - step_idx - 1], device=GPUS)
|
||||
alpha_prev = eval_alphas_prev[reversed_idx]
|
||||
ts = Tensor.full(bs, fill_value=timestep, dtype=dtypes.int, device="CPU")
|
||||
ts_ts = shard_tensor(ts.cat(ts))
|
||||
ts = shard_tensor(ts)
|
||||
sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[ts].reshape(bs, 1, 1, 1)
|
||||
sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[ts].reshape(bs, 1, 1, 1)
|
||||
x_x = shard_tensor(Tensor.stack(x.to("CPU"), x.to("CPU"), dim=1).reshape(-1, 4, 64, 64))
|
||||
x.assign(denoise_step(x, x_x, ts_ts, uc_c, sqrt_alphas_cumprod_t, sqrt_one_minus_alphas_cumprod_t, alpha_prev, unet, GPUS)).realize()
|
||||
return x
|
||||
|
||||
def decode_latents(latents:Tensor) -> Tensor: return vae_decode(shard_tensor(latents), first_stage, disable_beam=True)
|
||||
def generate_inception(imgs:Tensor) -> Tensor: return jit_inception(shard_tensor(imgs))[:,:,0,0]
|
||||
|
||||
def calc_clip_scores(batch:Tensor, batch_tokens:Tensor) -> Tensor:
|
||||
# Tensor.interpolate does not yet support bicubic, so we use PIL
|
||||
batch = (batch.to(GPUS[0]).permute(0,2,3,1) * 255).clip(0, 255).cast(dtypes.uint8).numpy()
|
||||
batch = [np.array(PIL.Image.fromarray(batch[i]).resize((224,224), PIL.Image.BICUBIC)) for i in range(bs)]
|
||||
batch = shard_tensor(Tensor(np.stack(batch, axis=0).transpose(0,3,1,2), device="CPU").realize())
|
||||
batch = batch.cast(dtypes.float) / 255
|
||||
batch = (batch - model.mean) / model.std
|
||||
batch = jit_clip(shard_tensor(batch_tokens), batch)
|
||||
return batch
|
||||
|
||||
callbacks = (embed_tokens, generate_latents, decode_latents, generate_inception, calc_clip_scores)
|
||||
|
||||
# save every forward pass output to disk; NOTE: this needs ~100 GB disk space because 30k images are large
|
||||
def stage_progress(stage_idx:int) -> int: return progress["end"].to("CPU")[stage_idx].item()
|
||||
if stage_progress(0) < len(eval_inputs):
|
||||
tokens = []
|
||||
for i in tqdm(range(0, len(eval_inputs), CONTEXT_BS)):
|
||||
subset = [cond_stage.tokenize(row["caption"], device="CPU") for row in eval_inputs[i: i+CONTEXT_BS]]
|
||||
tokens.append(Tensor.cat(*subset, dim=0).realize())
|
||||
progress["tokens"].assign(Tensor.cat(*tokens, dim=0).realize()).realize()
|
||||
progress["end"][0:1].assign(Tensor([len(eval_inputs)], dtype=dtypes.int)).realize()
|
||||
prev_stage = "tokens"
|
||||
tokens = progress["tokens"]
|
||||
|
||||
# wrapper code for every model
|
||||
for stage_idx, model, jit, bs, callback in zip(range(1,6), models, jits, all_bs, callbacks):
|
||||
stage = stages[stage_idx]
|
||||
if stage_progress(stage_idx) >= len(eval_inputs):
|
||||
prev_stage = stage
|
||||
continue # use cache
|
||||
t0 = time.perf_counter()
|
||||
print(f"starting eval with model: {model}")
|
||||
if stage_idx == 1: inputs = tokens
|
||||
elif stage_idx == 5: inputs = progress["imgs"]
|
||||
else: inputs = progress[prev_stage]
|
||||
|
||||
Tensor.realize(*[p.to_(GPUS) for p in get_parameters(model)])
|
||||
for batch_idx in tqdm(range(stage_progress(stage_idx), inputs.shape[0], bs)):
|
||||
t1 = time.perf_counter()
|
||||
batch, unpadded_bs = get_batch(inputs, batch_idx, bs)
|
||||
if isinstance(model, OpenClipEncoder): batch = callback(batch, get_batch(tokens, batch_idx, bs)[0].realize())
|
||||
else: batch = callback(batch)
|
||||
# to(GPUS[0]) is necessary for this to work, without that the result is still on GPUS, probably due to a bug
|
||||
batch = batch.to(GPUS[0]).to("CPU")[0:unpadded_bs].realize()
|
||||
progress[stage][batch_idx: batch_idx + bs].assign(batch).realize()
|
||||
# keep track of what our last output was, so we can resume from there if we crash in this loop
|
||||
progress["end"][stage_idx: stage_idx + 1].assign(Tensor([batch_idx + bs], dtype=dtypes.int)).realize()
|
||||
print(f"model: {model}, batch_idx: {batch_idx}, elapsed: {(time.perf_counter() - t1):.2f}")
|
||||
del batch
|
||||
|
||||
jit.reset()
|
||||
Tensor.realize(*[p.to_("CPU") for p in get_parameters(model)])
|
||||
print(f"done with model: {model}, elapsed: {(time.perf_counter() - t0):.2f}")
|
||||
prev_stage = stage
|
||||
|
||||
inception_stats_fn = str(DATADIR / "coco2014" / "val2014_30k_stats.npz")
|
||||
fid_score = inception.compute_score(progress["inception"].to("CPU"), inception_stats_fn)
|
||||
clip_score = progress["clip"].to(GPUS[0]).mean().item()
|
||||
for name in disk_tensor_names:
|
||||
Path(f"{EVAL_CKPT_DIR}/{name}.bytes").unlink(missing_ok=True)
|
||||
|
||||
if EVAL_SAMPLES and BEAM:
|
||||
print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
|
||||
sys.exit() # Don't eval additional models; we don't care about clip/fid scores when running BEAM on eval sample subset
|
||||
|
||||
return clip_score, fid_score
|
||||
|
||||
# evaluate checkpoints in reverse chronological order
|
||||
for ckpt_iteration, p in sorted(eval_queue, reverse=True):
|
||||
unet_ckpt = safe_load(p)
|
||||
load_state_dict(unet, unet_ckpt)
|
||||
clip_score, fid_score = eval_unet(eval_inputs, unet, model.cond_stage_model, model.first_stage_model, inception, clip_encoder)
|
||||
converged = True if clip_score >= 0.15 and fid_score <= 90 else False
|
||||
print(f"eval results for {EVAL_CKPT_DIR}/{p.name}: clip={clip_score}, fid={fid_score}, converged={converged}")
|
||||
if WANDB:
|
||||
wandb.log({"eval/ckpt_iteration": ckpt_iteration, "eval/clip_score": clip_score, "eval/fid_score": fid_score})
|
||||
if converged and STOP_IF_CONVERGED:
|
||||
print(f"Convergence detected, exiting early before evaluating other checkpoints due to STOP_IF_CONVERGED={STOP_IF_CONVERGED}")
|
||||
sys.exit()
|
||||
|
||||
# for testing
|
||||
return clip_score, fid_score, ckpt_iteration
|
||||
|
||||
if __name__ == "__main__":
|
||||
# inference only
|
||||
Tensor.training = False
|
||||
|
||||
@@ -1493,6 +1493,144 @@ def train_llama3():
|
||||
safe_save(get_state_dict(model), fn)
|
||||
break
|
||||
|
||||
def train_stable_diffusion():
|
||||
from extra.models.unet import UNetModel
|
||||
from examples.mlperf.dataloader import batch_load_train_stable_diffusion
|
||||
from examples.mlperf.lr_schedulers import LambdaLR, LambdaLinearScheduler
|
||||
from examples.mlperf.initializers import init_stable_diffusion
|
||||
from examples.mlperf.helpers import get_training_state
|
||||
import numpy as np
|
||||
|
||||
config = {}
|
||||
GPUS = config["GPUS"] = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 1))]
|
||||
seed = config["seed"] = getenv("SEED", 12345)
|
||||
# ** hyperparameters **
|
||||
BS = config["BS"] = getenv("BS", 1 * len(GPUS))
|
||||
BASE_LR = config["LEARNING_RATE"] = getenv("LEARNING_RATE", 2.5e-7)
|
||||
# https://github.com/mlcommons/training_policies/blob/cfa99da479b8d5931f7a3c67612d021dfb47510a/training_rules.adoc#benchmark_specific_rules
|
||||
# "Checkpoint must be collected every 512,000 images. CEIL(512000 / global_batch_size) if 512000 is not divisible by GBS."
|
||||
# NOTE: It's inferred that "steps" is the unit for the output of the CEIL formula, based on all other cases of CEIL in the rules
|
||||
CKPT_STEP_INTERVAL = config["CKPT_STEP_INTERVAL"] = getenv("CKPT_STEP_INTERVAL", math.ceil(512_000 / BS))
|
||||
CKPTDIR = config["CKPTDIR"] = Path(getenv("CKPTDIR", "./checkpoints"))
|
||||
DATADIR = config["DATADIR"] = Path(getenv("DATADIR", "./datasets"))
|
||||
UNET_CKPTDIR = config["UNET_CKPTDIR"] = Path(getenv("UNET_CKPTDIR", "./checkpoints"))
|
||||
TOTAL_CKPTS = config["TOTAL_CKPTS"] = getenv("TOTAL_CKPTS", 0)
|
||||
|
||||
print(f"training on {GPUS}")
|
||||
lr = BS * BASE_LR
|
||||
print(f"BS={BS}, BASE_LR={BASE_LR}, lr={lr}")
|
||||
print(f"CKPT_STEP_INTERVAL = {CKPT_STEP_INTERVAL}")
|
||||
for x in GPUS: Device[x]
|
||||
if (WANDB := getenv("WANDB", "")):
|
||||
import wandb
|
||||
wandb.init(config=config, project="MLPerf-Stable-Diffusion")
|
||||
|
||||
Tensor.manual_seed(seed) # seed for weight initialization
|
||||
model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod = init_stable_diffusion("v2-mlperf-train", CKPTDIR / "sd" / "512-base-ema.ckpt", GPUS)
|
||||
|
||||
optimizer = AdamW(get_parameters(unet))
|
||||
lambda_lr_callback = LambdaLinearScheduler(1000, 1.0, 1.0, 1e-06, 10000000000000).schedule
|
||||
lr_scheduler = LambdaLR(optimizer, Tensor(lr, dtype=dtypes.float, device=optimizer.device), lambda_lr_callback)
|
||||
|
||||
@TinyJit
|
||||
def train_step(mean:Tensor, logvar:Tensor, tokens:Tensor, unet:UNetModel, optimizer:LAMB, lr_scheduler:LambdaLR) -> Tensor:
|
||||
optimizer.zero_grad()
|
||||
|
||||
timestep = Tensor.randint(BS, low=0, high=model.alphas_cumprod.shape[0], dtype=dtypes.int, device=GPUS[0])
|
||||
latent_randn = Tensor.randn(*mean.shape, device=GPUS[0])
|
||||
noise = Tensor.randn(*mean.shape, device=GPUS[0])
|
||||
for t in (mean, logvar, tokens, timestep, latent_randn, noise):
|
||||
t.shard_(GPUS, axis=0)
|
||||
|
||||
std = Tensor.exp(0.5 * logvar.clamp(-30.0, 20.0))
|
||||
latent = (mean + std * latent_randn) * 0.18215
|
||||
|
||||
sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[timestep].reshape(timestep.shape[0], 1, 1, 1)
|
||||
sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[timestep].reshape(timestep.shape[0], 1, 1, 1)
|
||||
latent_with_noise = sqrt_alphas_cumprod_t * latent + sqrt_one_minus_alphas_cumprod_t * noise
|
||||
v_true = sqrt_alphas_cumprod_t * noise - sqrt_one_minus_alphas_cumprod_t * latent
|
||||
|
||||
context = model.cond_stage_model.embed_tokens(tokens)
|
||||
|
||||
out = unet(latent_with_noise, timestep, context)
|
||||
loss = ((out - v_true) ** 2).mean()
|
||||
del mean, logvar, std, latent, noise, sqrt_alphas_cumprod_t, sqrt_one_minus_alphas_cumprod_t
|
||||
del out, v_true, context, latent_randn, tokens, timestep
|
||||
loss.backward()
|
||||
|
||||
optimizer.step()
|
||||
lr_scheduler.step()
|
||||
loss, out_lr = loss.detach().to("CPU"), optimizer.lr.to("CPU")
|
||||
Tensor.realize(loss, out_lr)
|
||||
return loss, out_lr
|
||||
|
||||
# checkpointing takes ~9 minutes without this, and ~1 minute with this
|
||||
@TinyJit
|
||||
def ckpt_to_cpu():
|
||||
ckpt = get_training_state(unet, optimizer, lr_scheduler)
|
||||
# move to CPU first so more GPU bufs aren't created (can trigger OOM)
|
||||
for k,v in ckpt.items(): ckpt[k] = v.detach().to("CPU")
|
||||
Tensor.realize(*[v for v in ckpt.values()])
|
||||
for k,v in ckpt.items(): ckpt[k] = v.cast(v.dtype.base).contiguous()
|
||||
Tensor.realize(*[v for v in ckpt.values()])
|
||||
return ckpt
|
||||
|
||||
# training loop
|
||||
dl = batch_load_train_stable_diffusion(f'{DATADIR}/laion-400m/webdataset-moments-filtered/{{00000..00831}}.tar', BS)
|
||||
# for tests
|
||||
saved_checkpoints = []
|
||||
|
||||
train_start_time = time.perf_counter()
|
||||
t0 = t6 = time.perf_counter()
|
||||
for i, batch in enumerate(dl, start=1):
|
||||
loop_time = time.perf_counter() - t0
|
||||
t0 = time.perf_counter()
|
||||
dl_time = t0 - t6
|
||||
GlobalCounters.reset()
|
||||
|
||||
mean, logvar = np.split(np.concatenate(batch["npy"], axis=0), 2, axis=1)
|
||||
mean, logvar = Tensor(mean, dtype=dtypes.float32, device="CPU"), Tensor(logvar, dtype=dtypes.float32, device="CPU")
|
||||
tokens = []
|
||||
for text in batch['txt']: tokens += model.cond_stage_model.tokenizer.encode(text, pad_with_zeros=True)
|
||||
tokens = Tensor(tokens, dtype=dtypes.int32, device="CPU").reshape(-1, 77)
|
||||
|
||||
t1 = time.perf_counter()
|
||||
loss, lr = train_step(mean, logvar, tokens, unet, optimizer, lr_scheduler)
|
||||
loss_item, lr_item = loss.item(), lr.item()
|
||||
t2 = time.perf_counter()
|
||||
|
||||
if i == 3:
|
||||
for _ in range(3): ckpt_to_cpu() # do this at the beginning of run to prevent OOM surprises when checkpointing
|
||||
print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
|
||||
|
||||
total_train_time = time.perf_counter() - train_start_time
|
||||
if WANDB:
|
||||
wandb.log({"train/loss": loss_item, "train/lr": lr_item, "train/loop_time_prev": loop_time, "train/dl_time": dl_time, "train/step": i,
|
||||
"train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (t2-t1), "train/input_prep_time": t1-t0,
|
||||
"train/train_step_time": t2-t1, "train/total_time": total_train_time})
|
||||
|
||||
if i == 1 and wandb.run is not None:
|
||||
with open(f"{UNET_CKPTDIR}/wandb_run_id_{wandb.run.id}", "w") as f:
|
||||
f.write(f"wandb.run.id = {wandb.run.id}")
|
||||
|
||||
if i % CKPT_STEP_INTERVAL == 0:
|
||||
# https://github.com/mlcommons/training_policies/blob/cfa99da479b8d5931f7a3c67612d021dfb47510a/training_rules.adoc#benchmark_specific_rules
|
||||
# "evaluation is done offline, the time is not counted towards the submission time."
|
||||
fn = f"{UNET_CKPTDIR}/{i}.safetensors"
|
||||
print(f"saving unet checkpoint at {fn}")
|
||||
saved_checkpoints.append(fn)
|
||||
safe_save({k.replace("model.", ""):v for k,v in ckpt_to_cpu().items() if k.startswith("model.")}, fn)
|
||||
if TOTAL_CKPTS and i == TOTAL_CKPTS * CKPT_STEP_INTERVAL:
|
||||
print(f"ending run after {i} steps ({TOTAL_CKPTS} checkpoints collected)")
|
||||
return saved_checkpoints
|
||||
|
||||
t3 = time.perf_counter()
|
||||
print(f"""step {i}: {GlobalCounters.global_ops * 1e-9 / (t2-t1):9.2f} GFLOPS, mem_used: {GlobalCounters.mem_used / 1e9:.2f} GB,
|
||||
loop_time_prev: {loop_time:.2f}, dl_time: {dl_time:.2f}, input_prep_time: {t1-t0:.2f}, train_step_time: {t2-t1:.2f},
|
||||
t3-t2: {t3-t2:.4f}, loss:{loss_item:.5f}, lr:{lr_item:.3e}, total_train_time:{total_train_time:.2f}
|
||||
""")
|
||||
t6 = time.perf_counter()
|
||||
|
||||
if __name__ == "__main__":
|
||||
multiprocessing.set_start_method('spawn')
|
||||
|
||||
@@ -1501,7 +1639,7 @@ if __name__ == "__main__":
|
||||
else: bench_log_manager = contextlib.nullcontext()
|
||||
|
||||
with Tensor.train():
|
||||
for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn").split(","):
|
||||
for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn,stable_diffusion").split(","):
|
||||
nm = f"train_{m}"
|
||||
if nm in globals():
|
||||
print(f"training {m}")
|
||||
|
||||
+57
@@ -0,0 +1,57 @@
|
||||
#!/usr/bin/env bash
|
||||
# adapted from https://github.com/mlcommons/training/blob/4bdf5c8ed218ad76565a2ba1ac27c919ccc6d689/stable_diffusion/README.md
|
||||
|
||||
# setup dirs
|
||||
|
||||
DATA=/raid/datasets/stable_diffusion
|
||||
|
||||
LAION=$DATA/laion-400m/webdataset-moments-filtered
|
||||
COCO=$DATA/coco2014
|
||||
mkdir -p $LAION $COCO
|
||||
|
||||
CKPT=/raid/weights/stable_diffusion
|
||||
mkdir -p $CKPT/clip $CKPT/sd $CKPT/inception
|
||||
|
||||
# download data
|
||||
|
||||
# if rclone isn't installed system-wide / in your PATH, put the executable path in quotes below
|
||||
#RCLONE=""
|
||||
RCLONE="rclone"
|
||||
|
||||
## VAE-encoded image latents, from 6.1M image subset of laion-400m
|
||||
## about 1 TB for whole download
|
||||
$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
|
||||
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/ ${LAION} --include="*.tar" -P
|
||||
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/sha512sums.txt ${LAION} -P
|
||||
cd $LAION && grep -E '\.tar$' sha512sums.txt | sha512sum -c --quiet - && \
|
||||
echo "All .tar files verified" || { echo "Checksum failure when validating downloaded Laion moments"; exit 1; }
|
||||
|
||||
## prompts and FID statistics from 30k image subset of coco2014
|
||||
## 33 MB
|
||||
$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
|
||||
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k.tsv ${COCO} -P
|
||||
|
||||
$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
|
||||
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k_stats.npz ${COCO} -P
|
||||
|
||||
# download checkpoints
|
||||
|
||||
## clip (needed for text and vision encoders for validation)
|
||||
CLIP_WEIGHTS_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin"
|
||||
CLIP_WEIGHTS_SHA256="9a78ef8e8c73fd0df621682e7a8e8eb36c6916cb3c16b291a082ecd52ab79cc4"
|
||||
CLIP_CONFIG_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/raw/main/open_clip_config.json"
|
||||
wget -N -P ${CKPT}/clip ${CLIP_WEIGHTS_URL}
|
||||
wget -N -P ${CKPT}/clip ${CLIP_CONFIG_URL}
|
||||
echo "${CLIP_WEIGHTS_SHA256} ${CKPT}/clip/open_clip_pytorch_model.bin" | sha256sum -c
|
||||
|
||||
## sd (needed for latent->image decoder for validation, also has clip text encoder for training)
|
||||
SD_WEIGHTS_URL='https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt'
|
||||
SD_WEIGHTS_SHA256="d635794c1fedfdfa261e065370bea59c651fc9bfa65dc6d67ad29e11869a1824"
|
||||
wget -N -P ${CKPT}/sd ${SD_WEIGHTS_URL}
|
||||
echo "${SD_WEIGHTS_SHA256} ${CKPT}/sd/512-base-ema.ckpt" | sha256sum -c
|
||||
|
||||
## inception (needed for validation)
|
||||
FID_WEIGHTS_URL='https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
|
||||
FID_WEIGHTS_SHA1="bd836944fd6db519dfd8d924aa457f5b3c8357ff"
|
||||
wget -N -P ${CKPT}/inception ${FID_WEIGHTS_URL}
|
||||
echo "${FID_WEIGHTS_SHA1} ${CKPT}/inception/pt_inception-2015-12-05-6726825d.pth" | sha1sum -c
|
||||
+72
@@ -0,0 +1,72 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
DATETIME=${2:-$(date "+%m%d%H%M")}
|
||||
LOGFILE="${HOME}/logs/sd_mi300x_${DATETIME}.log"
|
||||
# UNET_CKPTDIR must be set: training saves checkpoints to this path, then a separate eval process scans this path to know which checkpoints to eval
|
||||
export UNET_CKPTDIR="${HOME}/stable_diffusion/training_checkpoints/${DATETIME}"
|
||||
mkdir -p "${HOME}/logs" "$UNET_CKPTDIR"
|
||||
|
||||
# run this script in isolation when using the --bg flag
|
||||
if [[ "${1:-}" == "--bg" ]]; then
|
||||
echo "logging output to $LOGFILE"
|
||||
echo "saving UNet checkpoints to $UNET_CKPTDIR"
|
||||
script_path="$(readlink -f "${BASH_SOURCE[0]}")"
|
||||
nohup bash "$script_path" run "$DATETIME" >"$LOGFILE" 2>&1 & disown $!
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# venv management
|
||||
if [[ -d .venv-sd-mlperf ]]; then
|
||||
. .venv-sd-mlperf/bin/activate
|
||||
else
|
||||
python3 -m venv .venv-sd-mlperf && . .venv-sd-mlperf/bin/activate
|
||||
pip install --index-url https://download.pytorch.org/whl/cpu torch && pip install tqdm numpy ftfy regex pillow scipy wandb webdataset
|
||||
fi
|
||||
pip list
|
||||
apt list --installed | grep amdgpu
|
||||
rocm-smi --version
|
||||
modinfo amdgpu | grep version
|
||||
|
||||
export BEAM=2 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 IGNORE_JIT_FIRST_BEAM=1 HCQDEV_WAIT_TIMEOUT_MS=300000
|
||||
export AMD_LLVM=0 # bf16 seems to require this
|
||||
export DATADIR="/raid/datasets/stable_diffusion"
|
||||
export CKPTDIR="/raid/weights/stable_diffusion"
|
||||
export EVAL_CKPT_DIR=$UNET_CKPTDIR
|
||||
export MODEL="stable_diffusion" PYTHONPATH="."
|
||||
export GPUS=8 BS=304
|
||||
export CONTEXT_BS=816 DENOISE_BS=600 DECODE_BS=384 INCEPTION_BS=560 CLIP_BS=240
|
||||
export WANDB=1
|
||||
export PARALLEL=4
|
||||
export PYTHONUNBUFFERED=1
|
||||
sudo rocm-smi -d 0 1 2 3 4 5 6 7 --setperfdeterminism 1500 || exit 1
|
||||
|
||||
# Retry BEAM search if script fails before BEAM COMPLETE is printed, but don't retry after that
|
||||
run_retry(){ local try=0 max=5 code tmp py pgid kids
|
||||
while :; do
|
||||
tmp=$(mktemp)
|
||||
setsid bash -c 'exec env "$@"' _ "$@" > >(tee -a "$LOGFILE" | tee "$tmp") 2>&1 &
|
||||
py=$!; pgid=$(ps -o pgid= -p "$py" | tr -d ' ')
|
||||
wait "$py"; code=$?
|
||||
[[ -n "$pgid" ]] && { kill -TERM -"$pgid" 2>/dev/null; sleep 1; kill -KILL -"$pgid" 2>/dev/null; }
|
||||
kids=$(pgrep -P "$py" || true)
|
||||
while [[ -n "$kids" ]]; do
|
||||
kill -TERM $kids 2>/dev/null; sleep 0.5
|
||||
kids=$(for k in $kids; do pgrep -P "$k" || true; done)
|
||||
done
|
||||
grep -q 'BEAM COMPLETE' "$tmp" && { rm -f "$tmp"; return 1; }
|
||||
rm -f "$tmp"
|
||||
((code==0)) && return 0
|
||||
((try>=max)) && return 2
|
||||
((try++)); sleep 90; echo "try = ${try}"
|
||||
done
|
||||
}
|
||||
|
||||
# Power limiting to 400W is only needed if GPUs fall out of sync (causing 2.2x increased train time) at higher power, which has been observed at 450W
|
||||
sudo rocm-smi -d 0 1 2 3 4 5 6 7 --setpoweroverdrive 750 && \
|
||||
run_retry TOTAL_CKPTS=7 python3 examples/mlperf/model_train.py; (( $? == 2 )) && { echo "training failed before BEAM completion"; exit 2; }
|
||||
sleep 90
|
||||
|
||||
run_retry EVAL_SAMPLES=600 python3 examples/mlperf/model_eval.py; (( $? == 2 )) && { echo "eval failed before BEAM completion"; exit 2; }
|
||||
# Checkpoints will be evaluated in reverse chronological order, even if above training crashed early
|
||||
# STOP_IF_CONVERGED=1: Stop the eval after the first time convergence is detected; no more checkpoints will be evaluated after that.
|
||||
STOP_IF_CONVERGED=1 python3 examples/mlperf/model_eval.py
|
||||
@@ -10,7 +10,7 @@ from tinygrad.helpers import DEBUG, getenv
|
||||
from tinygrad.engine.realize import CompiledRunner
|
||||
|
||||
import onnx
|
||||
from tinygrad.frontend.onnx import OnnxRunner
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
|
||||
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"
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
import sys
|
||||
from tinygrad import Tensor, fetch, GlobalCounters, dtypes
|
||||
from tinygrad.uop.ops import UOp
|
||||
from tinygrad.frontend.onnx import OnnxRunner
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from tinygrad.schedule.kernelize import get_kernelize_map
|
||||
from tinygrad.schedule.rangeify import get_rangeify_map
|
||||
from tinygrad.helpers import RANGEIFY
|
||||
from tinygrad.engine.schedule import create_schedule_with_vars
|
||||
from tinygrad.engine.realize import run_schedule
|
||||
|
||||
@@ -33,7 +35,7 @@ if __name__ == "__main__":
|
||||
if not in_target_path[s]:
|
||||
independent_set[s] = None
|
||||
independent = UOp.sink(*independent_set.keys())
|
||||
kernelized = get_kernelize_map(independent)
|
||||
kernelized = (get_rangeify_map if RANGEIFY else get_kernelize_map)(independent)
|
||||
independent = independent.substitute(kernelized)
|
||||
schedule, var_vals = create_schedule_with_vars(independent)
|
||||
run_schedule(schedule)
|
||||
|
||||
@@ -27,7 +27,7 @@ class Model(nn.Module):
|
||||
|
||||
if __name__ == "__main__":
|
||||
if getenv("TINY_BACKEND"):
|
||||
import tinygrad.frontend.torch # noqa: F401
|
||||
import tinygrad.nn.torch # noqa: F401
|
||||
device = torch.device("tiny")
|
||||
else:
|
||||
device = torch.device({"METAL":"mps","NV":"cuda"}.get(Device.DEFAULT, "cpu"))
|
||||
|
||||
@@ -9,11 +9,13 @@ from typing import Dict, Any
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
from tinygrad import Device, GlobalCounters, dtypes, Tensor, TinyJit
|
||||
from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm
|
||||
from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten
|
||||
from tinygrad.nn import Conv2d, GroupNorm
|
||||
from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
|
||||
from extra.models.clip import Closed, Tokenizer
|
||||
from extra.models.clip import Closed, Tokenizer, FrozenOpenClipEmbedder
|
||||
from extra.models import unet, clip
|
||||
from extra.models.unet import UNetModel
|
||||
from examples.mlperf.initializers import AutocastLinear, AutocastConv2d, AutocastGroupNorm, AutocastLayerNorm, zero_module, attn_f32_softmax, gelu_erf
|
||||
from extra.bench_log import BenchEvent, WallTimeEvent
|
||||
|
||||
class AttnBlock:
|
||||
@@ -154,12 +156,46 @@ unet_params: Dict[str,Any] = {
|
||||
"use_linear": False,
|
||||
}
|
||||
|
||||
mlperf_params: Dict[str,Any] = {"adm_in_ch": None, "in_ch": 4, "out_ch": 4, "model_ch": 320, "attention_resolutions": [4, 2, 1], "num_res_blocks": 2,
|
||||
"channel_mult": [1, 2, 4, 4], "d_head": 64, "transformer_depth": [1, 1, 1, 1], "ctx_dim": 1024, "use_linear": True,
|
||||
"num_groups":16, "st_norm_eps":1e-6}
|
||||
|
||||
class StableDiffusion:
|
||||
def __init__(self):
|
||||
def __init__(self, version:str|None=None, pretrained:str|None=None):
|
||||
self.alphas_cumprod = get_alphas_cumprod()
|
||||
self.model = namedtuple("DiffusionModel", ["diffusion_model"])(diffusion_model = UNetModel(**unet_params))
|
||||
self.first_stage_model = AutoencoderKL()
|
||||
self.cond_stage_model = namedtuple("CondStageModel", ["transformer"])(transformer = namedtuple("Transformer", ["text_model"])(text_model = Closed.ClipTextTransformer()))
|
||||
if version != "v2-mlperf-train":
|
||||
self.first_stage_model = AutoencoderKL() # only needed for decoding generated latents to images; not needed in mlperf training from preprocessed moments
|
||||
|
||||
if not version:
|
||||
self.cond_stage_model = namedtuple("CondStageModel", ["transformer"])(transformer = namedtuple("Transformer", ["text_model"])(text_model = Closed.ClipTextTransformer()))
|
||||
unet_init_params = unet_params
|
||||
elif version in {"v2-mlperf-train", "v2-mlperf-eval"}:
|
||||
unet_init_params = mlperf_params
|
||||
clip.gelu = gelu_erf
|
||||
self.cond_stage_model = FrozenOpenClipEmbedder(**{"dims": 1024, "n_heads": 16, "layers": 24, "return_pooled": False, "ln_penultimate": True,
|
||||
"clip_tokenizer_version": "sd_mlperf_v5_0"})
|
||||
unet.Linear, unet.Conv2d, unet.GroupNorm, unet.LayerNorm = AutocastLinear, AutocastConv2d, AutocastGroupNorm, AutocastLayerNorm
|
||||
unet.attention, unet.gelu, unet.mixed_precision_dtype = attn_f32_softmax, gelu_erf, dtypes.bfloat16
|
||||
if pretrained:
|
||||
print("loading text encoder")
|
||||
weights: dict[str,Tensor] = {k.replace("cond_stage_model.", "", 1):v for k,v in torch_load(pretrained)["state_dict"].items() if k.startswith("cond_stage_model.")}
|
||||
weights["model.attn_mask"] = Tensor.full((77, 77), fill_value=float("-inf")).triu(1)
|
||||
load_state_dict(self.cond_stage_model, weights)
|
||||
# only the eval model needs the decoder
|
||||
if version == "v2-mlperf-eval":
|
||||
print("loading image latent encoder")
|
||||
weights = {k.replace("first_stage_model.", "", 1):v for k,v in torch_load(pretrained)["state_dict"].items() if k.startswith("first_stage_model.")}
|
||||
load_state_dict(self.first_stage_model, weights)
|
||||
|
||||
self.model = namedtuple("DiffusionModel", ["diffusion_model"])(diffusion_model = UNetModel(**unet_init_params))
|
||||
if version == "v2-mlperf-train":
|
||||
# the mlperf reference inits certain weights as zeroes
|
||||
for bb in flatten(self.model.diffusion_model.input_blocks) + self.model.diffusion_model.middle_block + flatten(self.model.diffusion_model.output_blocks):
|
||||
if isinstance(bb, unet.ResBlock):
|
||||
zero_module(bb.out_layers[3])
|
||||
elif isinstance(bb, unet.SpatialTransformer):
|
||||
zero_module(bb.proj_out)
|
||||
zero_module(self.model.diffusion_model.out[2])
|
||||
|
||||
def get_x_prev_and_pred_x0(self, x, e_t, a_t, a_prev):
|
||||
temperature = 1
|
||||
|
||||
+1
-1
@@ -109,7 +109,7 @@ class TextDecoder:
|
||||
|
||||
def forward(self, x:Tensor, pos:Union[Variable, Literal[0]], encoded_audio:Tensor):
|
||||
seqlen = x.shape[-1]
|
||||
x = self.token_embedding(x) + self.positional_embedding.shrink(((pos, pos+seqlen), None, None))
|
||||
x = self.token_embedding(x) + self.positional_embedding.shrink(((pos, pos+seqlen), None))
|
||||
for block in self.blocks: x = block(x, xa=encoded_audio, mask=self.mask, len=pos)
|
||||
return self.output_tok(x)
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
import os
|
||||
from ultralytics import YOLO
|
||||
from pathlib import Path
|
||||
from tinygrad.frontend.onnx import OnnxRunner
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from extra.onnx_helpers import get_example_inputs
|
||||
|
||||
os.chdir("/tmp")
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import onnx, yaml, tempfile, time, argparse, json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from tinygrad.frontend.onnx import OnnxRunner
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from extra.onnx_helpers import validate, get_example_inputs
|
||||
from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry
|
||||
|
||||
|
||||
+32
-15
@@ -9,6 +9,9 @@ from PIL import Image
|
||||
import numpy as np
|
||||
import re, gzip
|
||||
|
||||
# Allow for monkeypatching for mlperf.
|
||||
gelu = Tensor.gelu
|
||||
|
||||
@lru_cache()
|
||||
def default_bpe():
|
||||
# Clip tokenizer, taken from https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py (MIT license)
|
||||
@@ -53,8 +56,8 @@ class Tokenizer:
|
||||
cs = [chr(n) for n in cs]
|
||||
return dict(zip(bs, cs))
|
||||
class ClipTokenizer:
|
||||
def __init__(self):
|
||||
self.byte_encoder = Tokenizer.bytes_to_unicode()
|
||||
def __init__(self, version=None):
|
||||
self.byte_encoder, self.version = Tokenizer.bytes_to_unicode(), version
|
||||
merges = gzip.open(default_bpe()).read().decode("utf-8").split('\n')
|
||||
merges = merges[1:49152-256-2+1]
|
||||
merges = [tuple(merge.split()) for merge in merges]
|
||||
@@ -62,11 +65,17 @@ class Tokenizer:
|
||||
vocab = vocab + [v+'</w>' for v in vocab]
|
||||
for merge in merges:
|
||||
vocab.append(''.join(merge))
|
||||
vocab.extend(['<|startoftext|>', '<|endoftext|>'])
|
||||
if self.version == "sd_mlperf_v5_0":
|
||||
import regex
|
||||
vocab.extend(['<start_of_text>', '<end_of_text>'])
|
||||
self.cache = {'<start_of_text>': '<start_of_text>', '<end_of_text>': '<end_of_text>'}
|
||||
self.pat = regex.compile(r"""<start_of_text>|<end_of_text>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", regex.IGNORECASE)
|
||||
else:
|
||||
vocab.extend(['<|startoftext|>', '<|endoftext|>'])
|
||||
self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
|
||||
self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[^\s]+""", re.IGNORECASE)
|
||||
self.encoder = dict(zip(vocab, range(len(vocab))))
|
||||
self.bpe_ranks = dict(zip(merges, range(len(merges))))
|
||||
self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
|
||||
self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[^\s]+""", re.IGNORECASE)
|
||||
|
||||
def bpe(self, token):
|
||||
if token in self.cache:
|
||||
@@ -110,8 +119,17 @@ class Tokenizer:
|
||||
|
||||
def encode(self, text:str, pad_with_zeros:bool=False) -> List[int]:
|
||||
bpe_tokens: List[int] = []
|
||||
text = Tokenizer.whitespace_clean(text.strip()).lower()
|
||||
for token in re.findall(self.pat, text):
|
||||
if self.version == "sd_mlperf_v5_0":
|
||||
import regex, ftfy, html
|
||||
text = ftfy.fix_text(text)
|
||||
text = html.unescape(html.unescape(text)).strip()
|
||||
text = Tokenizer.whitespace_clean(text).lower()
|
||||
re_module = regex
|
||||
else:
|
||||
text = Tokenizer.whitespace_clean(text.strip()).lower()
|
||||
re_module = re
|
||||
|
||||
for token in re_module.findall(self.pat, text):
|
||||
token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
|
||||
bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
|
||||
# Truncation, keeping two slots for start and end tokens.
|
||||
@@ -252,10 +270,8 @@ class Open:
|
||||
q,k,v = [y.reshape(T, B*self.n_heads, self.d_head).transpose(0, 1).reshape(B, self.n_heads, T, self.d_head) for y in proj.chunk(3)]
|
||||
|
||||
attn_output = Tensor.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
|
||||
attn_output = attn_output.permute(2, 0, 1, 3).reshape(T*B, C)
|
||||
|
||||
attn_output = attn_output.permute(2, 0, 1, 3).reshape(T, B, C)
|
||||
attn_output = self.out_proj(attn_output)
|
||||
attn_output = attn_output.reshape(T, B, C)
|
||||
|
||||
return attn_output
|
||||
|
||||
@@ -263,9 +279,10 @@ class Open:
|
||||
def __init__(self, dims, hidden_dims):
|
||||
self.c_fc = Linear(dims, hidden_dims)
|
||||
self.c_proj = Linear(hidden_dims, dims)
|
||||
self.gelu = gelu
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
return x.sequential([self.c_fc, Tensor.gelu, self.c_proj])
|
||||
return x.sequential([self.c_fc, self.gelu, self.c_proj])
|
||||
|
||||
# https://github.com/mlfoundations/open_clip/blob/58e4e39aaabc6040839b0d2a7e8bf20979e4558a/src/open_clip/transformer.py#L210
|
||||
class ResidualAttentionBlock:
|
||||
@@ -350,15 +367,15 @@ class Open:
|
||||
# https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/encoders/modules.py#L396
|
||||
# https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/encoders/modules.py#L498
|
||||
class FrozenOpenClipEmbedder(Embedder):
|
||||
def __init__(self, dims:int, n_heads:int, layers:int, return_pooled:bool, ln_penultimate:bool=False):
|
||||
self.tokenizer = Tokenizer.ClipTokenizer()
|
||||
def __init__(self, dims:int, n_heads:int, layers:int, return_pooled:bool, ln_penultimate:bool=False, clip_tokenizer_version=None):
|
||||
self.tokenizer = Tokenizer.ClipTokenizer(version=clip_tokenizer_version)
|
||||
self.model = Open.ClipTextTransformer(dims, n_heads, layers)
|
||||
self.return_pooled = return_pooled
|
||||
self.input_key = "txt"
|
||||
self.ln_penultimate = ln_penultimate
|
||||
|
||||
def tokenize(self, text:str, device:Optional[str]=None) -> Tensor:
|
||||
return Tensor(self.tokenizer.encode(text, pad_with_zeros=True), dtype=dtypes.int64, device=device).reshape(1,-1)
|
||||
return Tensor(self.tokenizer.encode(text, pad_with_zeros=True), dtype=dtypes.int32, device=device).reshape(1,-1)
|
||||
|
||||
def text_transformer_forward(self, x:Tensor, attn_mask:Optional[Tensor]=None):
|
||||
for r in self.model.transformer.resblocks:
|
||||
@@ -449,7 +466,7 @@ class OpenClipEncoder:
|
||||
x = x + self.positional_embedding
|
||||
x = self.transformer(x, attn_mask=self.attn_mask)
|
||||
x = self.ln_final(x)
|
||||
x = x[:, tokens.argmax(axis=-1)]
|
||||
x = x[Tensor.arange(x.shape[0], device=x.device), tokens.argmax(axis=-1)]
|
||||
x = x @ self.text_projection
|
||||
return x
|
||||
|
||||
|
||||
@@ -270,8 +270,10 @@ class FidInceptionV3:
|
||||
self.Mixed_7b = inception.Mixed_7b
|
||||
self.Mixed_7c = inception.Mixed_7c
|
||||
|
||||
def load_from_pretrained(self):
|
||||
state_dict = torch_load(str(fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")))
|
||||
def load_from_pretrained(self, path=None):
|
||||
if path is None:
|
||||
path = fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")
|
||||
state_dict = torch_load(str(path))
|
||||
for k,v in state_dict.items():
|
||||
if k.endswith(".num_batches_tracked"):
|
||||
state_dict[k] = v.reshape(1)
|
||||
|
||||
+35
-27
@@ -1,21 +1,24 @@
|
||||
from tinygrad import Tensor, dtypes
|
||||
from tinygrad.nn import Linear, Conv2d, GroupNorm, LayerNorm
|
||||
from tinygrad import Tensor, dtypes, nn
|
||||
from tinygrad.device import is_dtype_supported
|
||||
from typing import Optional, Union, List, Any, Tuple
|
||||
from typing import Optional, Union, List, Any, Tuple, Callable
|
||||
import math
|
||||
|
||||
# allow for monkeypatching
|
||||
Linear, Conv2d, GroupNorm, LayerNorm = nn.Linear, nn.Conv2d, nn.GroupNorm, nn.LayerNorm
|
||||
attention, gelu, mixed_precision_dtype = Tensor.scaled_dot_product_attention, Tensor.gelu, dtypes.float16
|
||||
|
||||
# https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/diffusionmodules/util.py#L207
|
||||
def timestep_embedding(timesteps:Tensor, dim:int, max_period=10000):
|
||||
half = dim // 2
|
||||
freqs = (-math.log(max_period) * Tensor.arange(half, device=timesteps.device) / half).exp()
|
||||
args = timesteps.unsqueeze(1) * freqs.unsqueeze(0)
|
||||
out = Tensor.cat(args.cos(), args.sin(), dim=-1)
|
||||
return out.cast(dtypes.float16) if is_dtype_supported(dtypes.float16) else out
|
||||
return out.cast(mixed_precision_dtype) if is_dtype_supported(mixed_precision_dtype) else out
|
||||
|
||||
class ResBlock:
|
||||
def __init__(self, channels:int, emb_channels:int, out_channels:int):
|
||||
def __init__(self, channels:int, emb_channels:int, out_channels:int, num_groups:int=32):
|
||||
self.in_layers = [
|
||||
GroupNorm(32, channels),
|
||||
GroupNorm(num_groups, channels),
|
||||
Tensor.silu,
|
||||
Conv2d(channels, out_channels, 3, padding=1),
|
||||
]
|
||||
@@ -24,7 +27,7 @@ class ResBlock:
|
||||
Linear(emb_channels, out_channels),
|
||||
]
|
||||
self.out_layers = [
|
||||
GroupNorm(32, out_channels),
|
||||
GroupNorm(num_groups, out_channels),
|
||||
Tensor.silu,
|
||||
lambda x: x, # needed for weights loading code to work
|
||||
Conv2d(out_channels, out_channels, 3, padding=1),
|
||||
@@ -45,35 +48,37 @@ class CrossAttention:
|
||||
self.to_v = Linear(ctx_dim, n_heads*d_head, bias=False)
|
||||
self.num_heads = n_heads
|
||||
self.head_size = d_head
|
||||
self.attn = attention
|
||||
self.to_out = [Linear(n_heads*d_head, query_dim)]
|
||||
|
||||
def __call__(self, x:Tensor, ctx:Optional[Tensor]=None) -> Tensor:
|
||||
ctx = x if ctx is None else ctx
|
||||
q,k,v = self.to_q(x), self.to_k(ctx), self.to_v(ctx)
|
||||
q,k,v = [y.reshape(x.shape[0], -1, self.num_heads, self.head_size).transpose(1,2) for y in (q,k,v)]
|
||||
attention = Tensor.scaled_dot_product_attention(q, k, v).transpose(1,2)
|
||||
attention = self.attn(q, k, v).transpose(1,2)
|
||||
h_ = attention.reshape(x.shape[0], -1, self.num_heads * self.head_size)
|
||||
return h_.sequential(self.to_out)
|
||||
|
||||
class GEGLU:
|
||||
def __init__(self, dim_in:int, dim_out:int):
|
||||
self.proj = Linear(dim_in, dim_out * 2)
|
||||
self.gelu = gelu
|
||||
self.dim_out = dim_out
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
x, gate = self.proj(x).chunk(2, dim=-1)
|
||||
return x * gate.gelu()
|
||||
return x * self.gelu(gate)
|
||||
|
||||
class FeedForward:
|
||||
def __init__(self, dim:int, mult:int=4):
|
||||
self.net = [
|
||||
self.net: tuple[GEGLU, Callable, nn.Linear] = (
|
||||
GEGLU(dim, dim*mult),
|
||||
lambda x: x, # needed for weights loading code to work
|
||||
Linear(dim*mult, dim)
|
||||
]
|
||||
)
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
return x.sequential(self.net)
|
||||
return x.sequential(list(self.net))
|
||||
|
||||
class BasicTransformerBlock:
|
||||
def __init__(self, dim:int, ctx_dim:int, n_heads:int, d_head:int):
|
||||
@@ -92,12 +97,13 @@ class BasicTransformerBlock:
|
||||
|
||||
# https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/attention.py#L619
|
||||
class SpatialTransformer:
|
||||
def __init__(self, channels:int, n_heads:int, d_head:int, ctx_dim:Union[int,List[int]], use_linear:bool, depth:int=1):
|
||||
def __init__(self, channels:int, n_heads:int, d_head:int, ctx_dim:Union[int,List[int]], use_linear:bool, depth:int=1,
|
||||
norm_eps:float=1e-5):
|
||||
if isinstance(ctx_dim, int):
|
||||
ctx_dim = [ctx_dim]*depth
|
||||
else:
|
||||
assert isinstance(ctx_dim, list) and depth == len(ctx_dim)
|
||||
self.norm = GroupNorm(32, channels)
|
||||
self.norm = GroupNorm(32, channels, eps=norm_eps)
|
||||
assert channels == n_heads * d_head
|
||||
self.proj_in = Linear(channels, channels) if use_linear else Conv2d(channels, channels, 1)
|
||||
self.transformer_blocks = [BasicTransformerBlock(channels, ctx_dim[d], n_heads, d_head) for d in range(depth)]
|
||||
@@ -134,7 +140,9 @@ class Upsample:
|
||||
|
||||
# https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/diffusionmodules/openaimodel.py#L472
|
||||
class UNetModel:
|
||||
def __init__(self, adm_in_ch:Optional[int], in_ch:int, out_ch:int, model_ch:int, attention_resolutions:List[int], num_res_blocks:int, channel_mult:List[int], transformer_depth:List[int], ctx_dim:Union[int,List[int]], use_linear:bool=False, d_head:Optional[int]=None, n_heads:Optional[int]=None):
|
||||
def __init__(self, adm_in_ch:Optional[int], in_ch:int, out_ch:int, model_ch:int, attention_resolutions:List[int], num_res_blocks:int,
|
||||
channel_mult:List[int], transformer_depth:List[int], ctx_dim:Union[int,List[int]], use_linear:bool=False, d_head:Optional[int]=None,
|
||||
n_heads:Optional[int]=None, num_groups:int=32, st_norm_eps:float=1e-5):
|
||||
self.model_ch = model_ch
|
||||
self.num_res_blocks = [num_res_blocks] * len(channel_mult)
|
||||
|
||||
@@ -174,12 +182,12 @@ class UNetModel:
|
||||
for idx, mult in enumerate(channel_mult):
|
||||
for _ in range(self.num_res_blocks[idx]):
|
||||
layers: List[Any] = [
|
||||
ResBlock(ch, time_embed_dim, model_ch*mult),
|
||||
ResBlock(ch, time_embed_dim, model_ch*mult, num_groups),
|
||||
]
|
||||
ch = mult * model_ch
|
||||
if ds in attention_resolutions:
|
||||
d_head, n_heads = get_d_and_n_heads(ch)
|
||||
layers.append(SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[idx]))
|
||||
layers.append(SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[idx], norm_eps=st_norm_eps))
|
||||
|
||||
self.input_blocks.append(layers)
|
||||
input_block_channels.append(ch)
|
||||
@@ -193,9 +201,9 @@ class UNetModel:
|
||||
|
||||
d_head, n_heads = get_d_and_n_heads(ch)
|
||||
self.middle_block: List = [
|
||||
ResBlock(ch, time_embed_dim, ch),
|
||||
SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[-1]),
|
||||
ResBlock(ch, time_embed_dim, ch),
|
||||
ResBlock(ch, time_embed_dim, ch, num_groups),
|
||||
SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[-1], norm_eps=st_norm_eps),
|
||||
ResBlock(ch, time_embed_dim, ch, num_groups),
|
||||
]
|
||||
|
||||
self.output_blocks = []
|
||||
@@ -203,13 +211,13 @@ class UNetModel:
|
||||
for i in range(self.num_res_blocks[idx] + 1):
|
||||
ich = input_block_channels.pop()
|
||||
layers = [
|
||||
ResBlock(ch + ich, time_embed_dim, model_ch*mult),
|
||||
ResBlock(ch + ich, time_embed_dim, model_ch*mult, num_groups),
|
||||
]
|
||||
ch = model_ch * mult
|
||||
|
||||
if ds in attention_resolutions:
|
||||
d_head, n_heads = get_d_and_n_heads(ch)
|
||||
layers.append(SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[idx]))
|
||||
layers.append(SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[idx], norm_eps=st_norm_eps))
|
||||
|
||||
if idx > 0 and i == self.num_res_blocks[idx]:
|
||||
layers.append(Upsample(ch))
|
||||
@@ -217,7 +225,7 @@ class UNetModel:
|
||||
self.output_blocks.append(layers)
|
||||
|
||||
self.out = [
|
||||
GroupNorm(32, ch),
|
||||
GroupNorm(num_groups, ch),
|
||||
Tensor.silu,
|
||||
Conv2d(model_ch, out_ch, 3, padding=1),
|
||||
]
|
||||
@@ -230,10 +238,10 @@ class UNetModel:
|
||||
assert y.shape[0] == x.shape[0]
|
||||
emb = emb + y.sequential(self.label_emb[0])
|
||||
|
||||
if is_dtype_supported(dtypes.float16):
|
||||
emb = emb.cast(dtypes.float16)
|
||||
ctx = ctx.cast(dtypes.float16)
|
||||
x = x .cast(dtypes.float16)
|
||||
if is_dtype_supported(mixed_precision_dtype):
|
||||
emb = emb.cast(mixed_precision_dtype)
|
||||
ctx = ctx.cast(mixed_precision_dtype)
|
||||
x = x .cast(mixed_precision_dtype)
|
||||
|
||||
def run(x:Tensor, bb) -> Tensor:
|
||||
if isinstance(bb, ResBlock): x = bb(x, emb)
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.tensor import _to_np_dtype
|
||||
from tinygrad.frontend.onnx import OnnxRunner, OnnxValue
|
||||
from tinygrad.nn.onnx import OnnxRunner, OnnxValue
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
|
||||
|
||||
@@ -50,7 +50,7 @@ class TestBeamSearch(unittest.TestCase):
|
||||
def test_variable_shrink_prime_number(self):
|
||||
v = Variable("v", 1, 400).bind(367)
|
||||
a = rand(400, 367)
|
||||
b = (a.shrink(((0,v), None))+1).reshape(367,367).realize()
|
||||
b = (a.shrink(((0,v), None))+1)[:367,:367].realize()
|
||||
np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)
|
||||
|
||||
def test_no_mutate_rawbuffers(self):
|
||||
|
||||
@@ -930,7 +930,7 @@ impl<'a> Thread<'a> {
|
||||
|
||||
let op = ((instr >> 16) & 0x3ff) as u32;
|
||||
match op {
|
||||
764 | 765 | 288 | 289 | 290 | 766 | 768 | 769 => {
|
||||
764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 => {
|
||||
let vdst = (instr & 0xff) as usize;
|
||||
let sdst = ((instr >> 8) & 0x7f) as usize;
|
||||
let f = |i: u32| -> usize { ((instr >> i) & 0x1ff) as usize };
|
||||
@@ -944,6 +944,16 @@ impl<'a> Thread<'a> {
|
||||
assert_eq!(clmp, 0);
|
||||
|
||||
let vcc = match op {
|
||||
767 => {
|
||||
let (s0, s1, s2): (u32, u32, u64) = (self.val(s0), self.val(s1), self.val(s2));
|
||||
let (mul_result, overflow_mul) = (s0 as i64).overflowing_mul(s1 as i64);
|
||||
let (ret, overflow_add) = mul_result.overflowing_add(s2 as i64);
|
||||
let overflowed = overflow_mul || overflow_add;
|
||||
if self.exec.read() {
|
||||
self.vec_reg.write64(vdst, ret as u64);
|
||||
}
|
||||
overflowed
|
||||
},
|
||||
766 => {
|
||||
let (s0, s1, s2): (u32, u32, u64) = (self.val(s0), self.val(s1), self.val(s2));
|
||||
let (mul_result, overflow_mul) = (s0 as u64).overflowing_mul(s1 as u64);
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
|
||||
Only supported on 7900XTX, requires either AM (`rmmod amdgpu`) or disabling power gating on AMD (`ppfeaturemask=0xffff3fff`, don't forget to rebuild initramfs)
|
||||
|
||||
SQTT is implemented on top of normal tinygrad PROFILE=1, `PROFILE=1 SQTT=1` to get profile pickle with sqtt data embedded in it.
|
||||
SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.
|
||||
|
||||
`SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.
|
||||
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
from tinygrad import Device, Tensor, Context
|
||||
from tinygrad.uop.ops import UOp, Ops, graph_rewrite, AxisType, PatternMatcher, UPat, pm_lower_index_dtype, GroupOp, KernelInfo
|
||||
from tinygrad.dtype import dtypes, AddrSpace
|
||||
from tinygrad.helpers import prod
|
||||
from tinygrad.schedule.rangeify import pm_mops
|
||||
from tinygrad.codegen.simplify import pm_flatten_range
|
||||
|
||||
TILE_DIM = 8
|
||||
N_BLOCK = 4
|
||||
K_BLOCK = 2
|
||||
M_BLOCK = 4
|
||||
|
||||
#M = N = K = 4096
|
||||
M = N = K = 1024
|
||||
|
||||
range_num = 0
|
||||
def rng(x, typ=AxisType.LOOP) -> UOp:
|
||||
global range_num
|
||||
range_num += 1
|
||||
return UOp.range(x, range_num-1, typ)
|
||||
|
||||
def glbl(nm, dtype, sz): return UOp(Ops.DEFINE_GLOBAL, dtype.ptr(prod(sz), AddrSpace.GLOBAL), arg=nm).reshape(sz)
|
||||
def rt(nm, dtype, sz): return UOp(Ops.DEFINE_REG, dtype.ptr(prod(sz), AddrSpace.REG), arg=nm).reshape(sz)
|
||||
|
||||
def zero(reg:UOp, *endrngs):
|
||||
rngs = [rng(s//TILE_DIM)*TILE_DIM for s in reg.shape]
|
||||
rngs = [x+rng(TILE_DIM) for x in rngs]
|
||||
|
||||
return reg[*rngs].store(UOp.const(reg.dtype.base, 0.0), *rngs, *endrngs, dtype=reg.dtype).reshape(reg.shape)
|
||||
|
||||
def load(reg:UOp, gl:UOp, *idxs):
|
||||
rngs = [rng(s//TILE_DIM)*TILE_DIM for s in reg.shape]
|
||||
rngs = [x+rng(TILE_DIM) for x in rngs]
|
||||
|
||||
grngs = [i*(r.vmax+1)+r for i,r in zip(idxs,rngs)]
|
||||
return reg[*rngs].store(gl[*grngs].load(), *rngs, dtype=reg.dtype).reshape(reg.shape)
|
||||
|
||||
def store(gl:UOp, reg:UOp, *idxs):
|
||||
rngs = [rng(s//TILE_DIM)*TILE_DIM for s in reg.shape]
|
||||
rngs = [x+rng(TILE_DIM) for x in rngs]
|
||||
|
||||
# TODO: why does this not have shape?
|
||||
#rngs = [rng(s) for s in (N_BLOCK*TILE_DIM, M_BLOCK*TILE_DIM)]
|
||||
grngs = [i*(r.vmax+1)+r for i,r in zip(idxs,rngs)]
|
||||
return gl[*grngs].store(reg[*rngs].load(), *rngs)
|
||||
|
||||
def mma_AB(outacc:UOp, a:UOp, b:UOp, *endrngs):
|
||||
assert a.shape[1] == b.shape[0]
|
||||
# meta::unroll_i_j_in_range -- split on TILE_DIM
|
||||
rngs = [rng(s//TILE_DIM)*TILE_DIM for s in outacc.shape]
|
||||
red = rng(a.shape[1]//TILE_DIM, AxisType.REDUCE)*TILE_DIM
|
||||
# meta::unroll_i_in_range -- split reduce on TILE_DIM
|
||||
rngs = [x+rng(TILE_DIM) for x in rngs]
|
||||
red = red + rng(TILE_DIM, AxisType.REDUCE)
|
||||
acc = outacc[*rngs].load(red) + a[rngs[0],red].load() * b[red,rngs[1]].load()
|
||||
return outacc[*rngs].store(acc, *rngs, red, *endrngs, dtype=outacc.dtype).reshape(outacc.shape)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# TODO: support string ranges
|
||||
tg_id_y = UOp.range(M // (M_BLOCK * TILE_DIM), -3, AxisType.GLOBAL if Device.DEFAULT != "CPU" else AxisType.LOOP)
|
||||
tg_id_x = UOp.range(N // (N_BLOCK * TILE_DIM), -2, AxisType.GLOBAL if Device.DEFAULT != "CPU" else AxisType.LOOP)
|
||||
|
||||
gl_d = glbl("gl0_d", dtypes.float, (N, M))
|
||||
gl_a = glbl("gl1_a", dtypes.float, (N, K))
|
||||
gl_b = glbl("gl2_b", dtypes.float, (K, M))
|
||||
|
||||
a_reg = rt("a_reg", dtypes.float, (N_BLOCK*TILE_DIM, K_BLOCK*TILE_DIM))
|
||||
b_reg = rt("b_reg", dtypes.float, (K_BLOCK*TILE_DIM, M_BLOCK*TILE_DIM))
|
||||
d_reg = rt("d_reg", dtypes.float, (N_BLOCK*TILE_DIM, M_BLOCK*TILE_DIM))
|
||||
d_reg = zero(d_reg, UOp(Ops.NOOP, src=(tg_id_y, tg_id_x)))
|
||||
|
||||
k = UOp.range(K // (K_BLOCK * TILE_DIM), -1, AxisType.REDUCE)
|
||||
a_reg = load(a_reg, gl_a, tg_id_y, k)
|
||||
b_reg = load(b_reg, gl_b, k, tg_id_x)
|
||||
d_reg = mma_AB(d_reg, a_reg, b_reg, k)
|
||||
sink = store(gl_d, d_reg, tg_id_y, tg_id_x).sink(arg=KernelInfo())
|
||||
|
||||
sink = graph_rewrite(sink, pm_mops+pm_flatten_range, name="pm_mops")
|
||||
|
||||
from tinygrad.codegen.gpudims import pm_add_gpudims
|
||||
sink = graph_rewrite(sink, pm_add_gpudims, ctx=Device.default.renderer, name="gpudims")
|
||||
|
||||
pm_lower_index_dtype_simple = PatternMatcher([
|
||||
(UPat(GroupOp.All, dtype=dtypes.index, name="x"), lambda x: x.replace(dtype=dtypes.int))
|
||||
])
|
||||
sink = graph_rewrite(sink, pm_lower_index_dtype_simple, name="index_dtype")
|
||||
|
||||
from tinygrad.codegen import rewrites_for_linearizer, apply_rewrites
|
||||
lin = apply_rewrites(sink, rewrites_for_linearizer)
|
||||
src = Device.default.renderer.render(lin.arg.lst)
|
||||
print(src)
|
||||
#exit(0)
|
||||
|
||||
from tinygrad.engine.realize import CompiledRunner, ExecItem
|
||||
from tinygrad.renderer import ProgramSpec
|
||||
|
||||
ps = ProgramSpec("test", src, Device.DEFAULT, sink, lin.arg.lst, [1,1,1], [1,1,1])
|
||||
run = CompiledRunner(ps)
|
||||
|
||||
a = Tensor.randn(N, N)
|
||||
b = Tensor.randn(N, N)
|
||||
c = Tensor.empty(N, N)
|
||||
Tensor.realize(a, b, c)
|
||||
|
||||
ei = ExecItem(run, [x.uop.buffer.ensure_allocated() for x in (c,a,b)])
|
||||
with Context(DEBUG=2):
|
||||
for i in range(5): ei.run()
|
||||
for i in range(5): ref = (a@b).realize()
|
||||
print((ref-c).mean().item())
|
||||
@@ -0,0 +1,74 @@
|
||||
# include directory copied from https://github.com/HazyResearch/ThunderMittens
|
||||
# https://hazyresearch.stanford.edu/blog/2024-11-28-tk-mlx
|
||||
|
||||
gemm = """
|
||||
#include <metal_stdlib>
|
||||
#include "include/tk.metal"
|
||||
using namespace mittens;
|
||||
|
||||
#define GEMM_PARAMS_DEF(T) \
|
||||
device T* D [[buffer(0)]], \
|
||||
device T* A [[buffer(1)]], \
|
||||
device T* B [[buffer(2)]], \
|
||||
const constant int &N [[buffer(3)]], \
|
||||
const constant int &K [[buffer(4)]], \
|
||||
const constant int &M [[buffer(5)]], \
|
||||
uint3 tg_id [[threadgroup_position_in_grid]], \
|
||||
uint simd_lane_id [[thread_index_in_simdgroup]]
|
||||
|
||||
template<typename T, unsigned N_BLOCK, unsigned K_BLOCK, unsigned M_BLOCK>
|
||||
kernel void matmul_naive(GEMM_PARAMS_DEF(T)) {
|
||||
using global_layout = gl<T, 1, 1, -1, -1>;
|
||||
global_layout gl_a(A, nullptr, nullptr, N, K);
|
||||
global_layout gl_b(B, nullptr, nullptr, K, M);
|
||||
global_layout gl_d(D, nullptr, nullptr, N, M);
|
||||
rt<T, N_BLOCK * TILE_DIM, K_BLOCK * TILE_DIM> a_reg;
|
||||
rt<T, K_BLOCK * TILE_DIM, M_BLOCK * TILE_DIM> b_reg;
|
||||
rt<float, N_BLOCK * TILE_DIM, M_BLOCK * TILE_DIM> d_reg;
|
||||
zero(d_reg);
|
||||
#pragma clang loop unroll(full)
|
||||
for (int k = 0; k < K / (K_BLOCK * TILE_DIM); k++) {
|
||||
load(a_reg, gl_a, {0, 0, (int)tg_id.y, k}, simd_lane_id);
|
||||
load(b_reg, gl_b, {0, 0, k, (int)tg_id.x}, simd_lane_id);
|
||||
mma_AB(d_reg, a_reg, b_reg, d_reg);
|
||||
}
|
||||
store(gl_d, d_reg, {0, 0, (int)tg_id.y, (int)tg_id.x}, simd_lane_id);
|
||||
}
|
||||
|
||||
#define instantiate_matmul_custom(type_name, T) \
|
||||
template [[host_name("matmul_custom_" #type_name)]] [[kernel]] \
|
||||
void matmul_naive<T, 4, 2, 4>(GEMM_PARAMS_DEF(T)); \
|
||||
|
||||
instantiate_matmul_custom(float32, float);
|
||||
"""
|
||||
|
||||
from tinygrad import Device, Tensor, Context
|
||||
|
||||
if __name__ == "__main__":
|
||||
device = Device["METAL"]
|
||||
lib = device.compiler.compile(gemm)
|
||||
prg = device.runtime("matmul_custom_float32", lib)
|
||||
|
||||
N = 4096
|
||||
a = Tensor.randn(N, N)
|
||||
b = Tensor.randn(N, N)
|
||||
c = Tensor.empty(N, N)
|
||||
Tensor.realize(a, b, c)
|
||||
|
||||
TILE_DIM = 8
|
||||
N_BLOCK = 4
|
||||
M_BLOCK = 4
|
||||
|
||||
gsz = (N // (M_BLOCK * TILE_DIM), N // (N_BLOCK * TILE_DIM), 1)
|
||||
for _ in range(5):
|
||||
et = prg(c.uop.buffer.ensure_allocated()._buf, a.uop.buffer._buf, b.uop.buffer._buf,
|
||||
global_size=gsz, local_size=(32,1,1), vals=(N, N, N), wait=True)
|
||||
print(f"{N*N*N*2/(et*1e9):2f} GFLOPS")
|
||||
|
||||
for _ in range(5):
|
||||
with Context(DEBUG=2):
|
||||
ref = (a@b).realize()
|
||||
|
||||
print((ref-c).mean().item())
|
||||
|
||||
|
||||
@@ -0,0 +1,392 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Basic operations on generic types.
|
||||
*/
|
||||
#pragma once
|
||||
#include "base_types.metal"
|
||||
#include <metal_math>
|
||||
|
||||
namespace mittens {
|
||||
/**
|
||||
* @namespace base_ops
|
||||
*
|
||||
* @brief A namespace for operations on basic data types.
|
||||
*/
|
||||
namespace base_ops {
|
||||
#define TEMPLATE_OPS_SINGLE(func_contents) \
|
||||
template<typename T> static METAL_FUNC T op(device const T &x) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &x) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &x) { func_contents }
|
||||
|
||||
#define TEMPLATE_OPS_OVERRIDE_SINGLE(T, op_name, func_contents) \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &x) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &x) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &x) { func_contents }
|
||||
|
||||
#define TEMPLATE_OPS_DOUBLE(func_contents) \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, device const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, threadgroup const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, thread const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, device const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, threadgroup const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, thread const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, device const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, threadgroup const T &b) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, thread const T &b) { func_contents }
|
||||
|
||||
#define TEMPLATE_OPS_OVERRIDE_DOUBLE(T, op_name, func_contents) \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, device const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, threadgroup const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, thread const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, device const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, threadgroup const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, thread const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, device const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, threadgroup const T &b) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, thread const T &b) { func_contents }
|
||||
|
||||
#define TEMPLATE_OPS_TRIPLE(func_contents) \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, device const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, device const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, device const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, threadgroup const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, threadgroup const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, thread const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, thread const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(device const T &a, thread const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, device const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, device const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, device const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, threadgroup const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, threadgroup const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, thread const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, thread const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &a, thread const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, device const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, device const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, device const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, threadgroup const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, threadgroup const T &b, thread const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, thread const T &b, device const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, thread const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &a, thread const T &b, thread const T &c) { func_contents }
|
||||
|
||||
#define TEMPLATE_OPS_OVERRIDE_TRIPLE(T, op_name, func_contents) \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, device const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, device const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, device const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, threadgroup const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, threadgroup const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, thread const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, thread const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(device const T &a, thread const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, device const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, device const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, device const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, threadgroup const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, threadgroup const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, thread const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, thread const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, thread const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, device const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, device const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, device const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, threadgroup const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, threadgroup const T &b, thread const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, thread const T &b, device const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, thread const T &b, threadgroup const T &c) { func_contents } \
|
||||
template<> METAL_FUNC T op_name::op<T>(thread const T &a, thread const T &b, thread const T &c) { func_contents }
|
||||
|
||||
|
||||
|
||||
/* ---------- 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> static METAL_FUNC 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> static METAL_FUNC 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> static METAL_FUNC 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> static METAL_FUNC 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_OPS_SINGLE(return metal::exp(x);)
|
||||
};
|
||||
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16, exp, return bf16(metal::exp((float)x));)
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, exp, return bf16_2(metal::exp(float2(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 METAL_FUNC T op(device const T &x) { return metal::exp2(x); } \
|
||||
template<typename T> static METAL_FUNC T op(threadgroup const T &x) { return metal::exp2(x); } \
|
||||
template<typename T> static METAL_FUNC T op(thread const T &x) { return metal::exp2(x); }
|
||||
};
|
||||
|
||||
//template<> METAL_FUNC bf16 exp2::op<bf16>(device const bf16 &x) { return bf16(metal::exp2(x)); } \
|
||||
//template<> METAL_FUNC bf16 exp2::op<bf16>(threadgroup const bf16 &x) { return bf16(metal::exp2(x)); } \
|
||||
//template<> METAL_FUNC bf16 exp2::op<bf16>(thread const bf16 &x) { return bf16(metal::exp2(x)); }
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16, exp2, return bf16(metal::exp2(x));)
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, exp2, return bf16_2(metal::exp2((float2)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_OPS_SINGLE(return metal::log(x);)
|
||||
};
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16, log, return bf16(metal::log(x));)
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, log, return bf16_2(metal::log((float2)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_OPS_SINGLE(return metal::abs(x);)
|
||||
};
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16 , abs, return bf16(metal::abs((float)x));)
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, abs, return bf16_2(metal::abs((float2)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_OPS_SINGLE(return max(x, base_types::constants<T>::zero());)
|
||||
};
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16 , relu, return bf16(metal::max((float)x, base_types::constants<float>::zero()));)
|
||||
TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, relu, return bf16_2(metal::max((float2)x, base_types::constants<float2>::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_OPS_SINGLE(return x;)
|
||||
};
|
||||
|
||||
/* ---------- 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_OPS_DOUBLE(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_OPS_DOUBLE(return 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_OPS_DOUBLE(return 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_OPS_DOUBLE(return 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_OPS_DOUBLE(return 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_OPS_DOUBLE(return metal::max(a,b);)
|
||||
};
|
||||
TEMPLATE_OPS_OVERRIDE_DOUBLE(bf16 , max, return (bf16)metal::max((float)a, (float)b);)
|
||||
TEMPLATE_OPS_OVERRIDE_DOUBLE(bf16_2, max, return (bf16_2)metal::max((float2)a, (float2)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_OPS_DOUBLE(return metal::min(a,b);)
|
||||
};
|
||||
TEMPLATE_OPS_OVERRIDE_DOUBLE(bf16 , min, return (bf16)metal::min((float)a, (float)b);)
|
||||
TEMPLATE_OPS_OVERRIDE_DOUBLE(bf16_2, min, return (bf16_2)metal::min((float2)a, (float2)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_OPS_TRIPLE(return sum::op<T>(mul::op<T>(a, b), c);)
|
||||
};
|
||||
|
||||
/**
|
||||
* @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_OPS_TRIPLE(return sum::op<T>(mul::op<T>(a, c), b);)
|
||||
};
|
||||
|
||||
#undef TEMPLATE_OPS_SINGLE
|
||||
#undef TEMPLATE_OPS_OVERRIDE_SINGLE
|
||||
#undef TEMPLATE_OPS_DOUBLE
|
||||
#undef TEMPLATE_OPS_OVERRIDE_DOUBLE
|
||||
#undef TEMPLATE_OPS_TRIPLE
|
||||
#undef TEMPLATE_OPS_OVERRIDE_TRIPLE
|
||||
} // base_ops
|
||||
} // mittens
|
||||
@@ -0,0 +1,321 @@
|
||||
|
||||
#pragma once
|
||||
|
||||
namespace mittens {
|
||||
|
||||
using bf16 = bfloat;
|
||||
using bf16_2 = bfloat2;
|
||||
using bf16_4 = bfloat4;
|
||||
//using half_2 = half2;
|
||||
|
||||
namespace ducks {
|
||||
namespace base_types {
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isT1() {
|
||||
return metal::is_same<typename T::dtype, float>::value ||
|
||||
metal::is_same<typename T::dtype, bf16 >::value ||
|
||||
metal::is_same<typename T::dtype, half>::value;
|
||||
}
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isT2() {
|
||||
return metal::is_same<typename T::dtype, float2>::value ||
|
||||
metal::is_same<typename T::dtype, bf16_2>::value ||
|
||||
metal::is_same<typename T::dtype, half2>::value;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isT1Type() {
|
||||
return metal::is_same<T, float>::value ||
|
||||
metal::is_same<T, bf16 >::value ||
|
||||
metal::is_same<T, half>::value;
|
||||
}
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isT2Type() {
|
||||
return metal::is_same<T, float2>::value ||
|
||||
metal::is_same<T, bf16_2>::value ||
|
||||
metal::is_same<T, half2>::value;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isT1Ptr() {
|
||||
return metal::is_same<T, device float*>::value ||
|
||||
metal::is_same<T, threadgroup float*>::value ||
|
||||
metal::is_same<T, thread float*>::value ||
|
||||
metal::is_same<T, device bf16*>::value ||
|
||||
metal::is_same<T, threadgroup bf16*>::value ||
|
||||
metal::is_same<T, thread bf16*>::value ||
|
||||
metal::is_same<T, device half*>::value ||
|
||||
metal::is_same<T, threadgroup half*>::value ||
|
||||
metal::is_same<T, thread half*>::value;
|
||||
}
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isT2Ptr() {
|
||||
return metal::is_same<T, device float2*>::value ||
|
||||
metal::is_same<T, threadgroup float2*>::value ||
|
||||
metal::is_same<T, thread float2*>::value ||
|
||||
metal::is_same<T, device bf16_2*>::value ||
|
||||
metal::is_same<T, threadgroup bf16_2*>::value ||
|
||||
metal::is_same<T, thread bf16_2*>::value ||
|
||||
metal::is_same<T, device half2*>::value ||
|
||||
metal::is_same<T, threadgroup half2*>::value ||
|
||||
metal::is_same<T, thread half2*>::value;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static METAL_FUNC constexpr const bool isTKType() { // good enough
|
||||
return !isT1Type<T>() && !isT2Type<T>() && !isT1Ptr<T>() && !isT2Ptr<T>();
|
||||
}
|
||||
|
||||
} // namespace base_types
|
||||
} // namespace ducks
|
||||
|
||||
/**
|
||||
* @namespace base_types
|
||||
*
|
||||
* @brief A namespace for Thundermittens 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 METAL_FUNC constexpr T zero() { return T{0}; }
|
||||
/**
|
||||
* @brief One
|
||||
* @return Constexpr one with type T
|
||||
*/
|
||||
static METAL_FUNC 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 METAL_FUNC 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 METAL_FUNC constexpr T neg_infty() { return T{-INFINITY}; }
|
||||
};
|
||||
template<> struct constants<float> {
|
||||
static METAL_FUNC constexpr float zero() { return 0.f; }
|
||||
static METAL_FUNC constexpr float one() { return 1.f; }
|
||||
static METAL_FUNC constexpr float pos_infty() { return INFINITY; }
|
||||
static METAL_FUNC constexpr float neg_infty() { return -INFINITY; }
|
||||
};
|
||||
template<> struct constants<float2> {
|
||||
static METAL_FUNC constexpr float2 zero() { return float2(0.f, 0.f); }
|
||||
static METAL_FUNC constexpr float2 one() { return float2(1.f, 1.f); }
|
||||
static METAL_FUNC constexpr float2 pos_infty() { return float2(constants<float>::pos_infty(), constants<float>::pos_infty()); }
|
||||
static METAL_FUNC constexpr float2 neg_infty() { return float2(constants<float>::neg_infty(), constants<float>::neg_infty()); }
|
||||
};
|
||||
template<> struct constants<bf16> {
|
||||
static METAL_FUNC constexpr bf16 zero() { return 0.bf; }
|
||||
static METAL_FUNC constexpr bf16 one() { return 1.bf; }
|
||||
static METAL_FUNC constexpr bf16 pos_infty() { return HUGE_VALBF; }
|
||||
static METAL_FUNC constexpr bf16 neg_infty() { return -HUGE_VALBF; }
|
||||
};
|
||||
template<> struct constants<bf16_2> {
|
||||
static METAL_FUNC constexpr bf16_2 zero() { return bf16_2(constants<bf16>::zero(), constants<bf16>::zero()); }
|
||||
static METAL_FUNC constexpr bf16_2 one() { return bf16_2(constants<bf16>::one(), constants<bf16>::one()); }
|
||||
static METAL_FUNC constexpr bf16_2 pos_infty() { return bf16_2(constants<bf16>::pos_infty(), constants<bf16>::pos_infty()); }
|
||||
static METAL_FUNC constexpr bf16_2 neg_infty() { return bf16_2(constants<bf16>::neg_infty(), constants<bf16>::neg_infty()); }
|
||||
};
|
||||
template<> struct constants<half> {
|
||||
static METAL_FUNC constexpr half zero() { return half(0.h); }
|
||||
static METAL_FUNC constexpr half one() { return half(1.h); }
|
||||
static METAL_FUNC constexpr half pos_infty() { return HUGE_VALH; }
|
||||
static METAL_FUNC constexpr half neg_infty() { return -HUGE_VALH; }
|
||||
};
|
||||
|
||||
template<> struct constants<half2> {
|
||||
static METAL_FUNC constexpr half2 zero() { return half2(constants<half>::zero(), constants<half>::zero()); }
|
||||
static METAL_FUNC constexpr half2 one() { return half2(constants<half>::one(), constants<half>::one()); }
|
||||
static METAL_FUNC constexpr half2 pos_infty() { return half2(constants<half>::pos_infty(), constants<half>::pos_infty()); }
|
||||
static METAL_FUNC constexpr half2 neg_infty() { return half2(constants<half>::neg_infty(), constants<half>::neg_infty()); }
|
||||
};
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* @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 METAL_FUNC 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 METAL_FUNC constexpr T pack(device const bf16 &i);
|
||||
// static METAL_FUNC constexpr T pack(threadgroup const bf16 &i);
|
||||
// static METAL_FUNC constexpr T pack(thread const bf16 &i);
|
||||
};
|
||||
|
||||
#define PACK_FUNCTIONS(T1, T2) \
|
||||
static METAL_FUNC constexpr T2 pack(device const T1 &i) { return T2{i, i}; } \
|
||||
static METAL_FUNC constexpr T2 pack(threadgroup const T1 &i) { return T2{i, i}; } \
|
||||
static METAL_FUNC constexpr T2 pack(thread const T1 &i) { return T2{i, i}; }
|
||||
|
||||
template<> struct packing<bf16> {
|
||||
static METAL_FUNC constexpr int num() { return 1; }
|
||||
using unpacked_type = bf16;
|
||||
using packed_type = bf16_2;
|
||||
using packed_four = bf16_4;
|
||||
PACK_FUNCTIONS(unpacked_type, packed_type)
|
||||
};
|
||||
template<> struct packing<half> {
|
||||
static METAL_FUNC constexpr int num() { return 1; }
|
||||
using unpacked_type = half;
|
||||
using packed_type = half2;
|
||||
using packed_four = half4;
|
||||
PACK_FUNCTIONS(unpacked_type, packed_type)
|
||||
};
|
||||
template<> struct packing<float> {
|
||||
static METAL_FUNC constexpr int num() { return 1; }
|
||||
using unpacked_type = float;
|
||||
using packed_type = float2;
|
||||
using packed_four = float4;
|
||||
|
||||
PACK_FUNCTIONS(unpacked_type, packed_type)
|
||||
};
|
||||
template<> struct packing<bf16_2> {
|
||||
static METAL_FUNC constexpr int num() { return 2; }
|
||||
using unpacked_type = bf16;
|
||||
using packed_type = bf16_2;
|
||||
using packed_four = bf16_4;
|
||||
PACK_FUNCTIONS(unpacked_type, packed_type)
|
||||
};
|
||||
template<> struct packing<half2> {
|
||||
static METAL_FUNC constexpr int num() { return 2; }
|
||||
using unpacked_type = half;
|
||||
using packed_type = half2;
|
||||
using packed_four = half4;
|
||||
PACK_FUNCTIONS(unpacked_type, packed_type)
|
||||
};
|
||||
template<> struct packing<float2> {
|
||||
static METAL_FUNC constexpr int num() { return 2; }
|
||||
using unpacked_type = float;
|
||||
using packed_type = float2;
|
||||
using packed_four = float4;
|
||||
PACK_FUNCTIONS(unpacked_type, packed_type)
|
||||
};
|
||||
template<> struct packing<int2> {
|
||||
static METAL_FUNC constexpr int num() { return 2; }
|
||||
};
|
||||
template<> struct packing<float4> {
|
||||
static METAL_FUNC constexpr int num() { return 4; }
|
||||
};
|
||||
template<> struct packing<int4> {
|
||||
static METAL_FUNC constexpr int num() { return 4; }
|
||||
};
|
||||
|
||||
|
||||
/**
|
||||
* @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 METAL_FUNC T convert(device const U & u) { return (T)u; }
|
||||
static METAL_FUNC T convert(threadgroup const U & u) { return (T)u; }
|
||||
static METAL_FUNC T convert(thread const U & u) { return (T)u; }
|
||||
};
|
||||
|
||||
template<> struct convertor<float, bf16> {
|
||||
// fptrunc float %_ to bfloat
|
||||
static METAL_FUNC float convert(device const bf16 & u) { return float(u);}
|
||||
static METAL_FUNC float convert(threadgroup const bf16 & u) { return float(u);}
|
||||
static METAL_FUNC float convert(thread const bf16 & u) { return float(u);}
|
||||
};
|
||||
template<> struct convertor<bf16, float> {
|
||||
// fpext bfloat %_ to float
|
||||
static METAL_FUNC bf16 convert(device const float & u) { return bf16(u); }
|
||||
static METAL_FUNC bf16 convert(threadgroup const float & u) { return bf16(u); }
|
||||
static METAL_FUNC bf16 convert(thread const float & u) { return bf16(u); }
|
||||
};
|
||||
template<> struct convertor<float2, bf16_2> {
|
||||
// tail call fast <2 x float> @air.convert.f.v2f32.f.v2bf16(<2 x bfloat> %_)
|
||||
static METAL_FUNC float2 convert(device const bf16_2 & u) { return float2(u); }
|
||||
static METAL_FUNC float2 convert(threadgroup const bf16_2 & u) { return float2(u); }
|
||||
static METAL_FUNC float2 convert(thread const bf16_2 & u) { return float2(u); }
|
||||
};
|
||||
template<> struct convertor<bf16_2, float2> {
|
||||
// tail call fast <2 x bfloat> @air.convert.f.v2bf16.f.v2f32(<2 x float> %_)
|
||||
static METAL_FUNC bf16_2 convert(device const float2 & u) { return bf16_2(u); }
|
||||
static METAL_FUNC bf16_2 convert(threadgroup const float2 & u) { return bf16_2(u); }
|
||||
static METAL_FUNC bf16_2 convert(thread const float2 & u) { return bf16_2(u); }
|
||||
};
|
||||
|
||||
template<> struct convertor<float, half> {
|
||||
// fptrunc float %_ to half
|
||||
static METAL_FUNC float convert(device const half & u) { return float(u); }
|
||||
static METAL_FUNC float convert(threadgroup const half & u) { return float(u); }
|
||||
static METAL_FUNC float convert(thread const half & u) { return float(u); }
|
||||
};
|
||||
template<> struct convertor<half, float> {
|
||||
//fpext half %_ to float
|
||||
static METAL_FUNC half convert(device const float & u) { return half(u); }
|
||||
static METAL_FUNC half convert(threadgroup const float & u) { return half(u); }
|
||||
static METAL_FUNC half convert(thread const float & u) { return half(u); }
|
||||
};
|
||||
template<> struct convertor<float2, half2> {
|
||||
// tail call fast <2 x float> @air.convert.f.v2f32.f.v2f16(<2 x half> %_)
|
||||
static METAL_FUNC float2 convert(device const half2 & u) { return float2(u); }
|
||||
static METAL_FUNC float2 convert(threadgroup const half2 & u) { return float2(u); }
|
||||
static METAL_FUNC float2 convert(thread const half2 & u) { return float2(u); }
|
||||
};
|
||||
template<> struct convertor<half2, float2> {
|
||||
// tail call fast <2 x half> @air.convert.f.v2f16.f.v2f32(<2 x float> %_)
|
||||
static METAL_FUNC half2 convert(device const float2 & u) { return half2(u); }
|
||||
static METAL_FUNC half2 convert(threadgroup const float2 & u) { return half2(u); }
|
||||
static METAL_FUNC half2 convert(thread const float2 & u) { return half2(u); }
|
||||
};
|
||||
template<> struct convertor<bf16, half> {
|
||||
static METAL_FUNC bf16 convert(device const half & u) { return bf16(u); }
|
||||
static METAL_FUNC bf16 convert(threadgroup const half & u) { return bf16(u); }
|
||||
static METAL_FUNC bf16 convert(thread const half & u) { return bf16(u); }
|
||||
};
|
||||
template<> struct convertor<half, bf16> {
|
||||
static METAL_FUNC half convert(device const bf16 & u) { return half(u); }
|
||||
static METAL_FUNC half convert(threadgroup const bf16 & u) { return half(u); }
|
||||
static METAL_FUNC half convert(thread const bf16 & u) { return half(u); }
|
||||
};
|
||||
template<> struct convertor<bf16_2, half2> {
|
||||
// tail call fast <2 x bfloat> @air.convert.f.v2bf16.f.v2f16(<2 x half> %_)
|
||||
static METAL_FUNC bf16_2 convert(device const half2 & u) { return bf16_2(u); }
|
||||
static METAL_FUNC bf16_2 convert(threadgroup const half2 & u) { return bf16_2(u); }
|
||||
static METAL_FUNC bf16_2 convert(thread const half2 & u) { return bf16_2(u); }
|
||||
};
|
||||
template<> struct convertor<half2, bf16_2> {
|
||||
// tail call fast <2 x half> @air.convert.f.v2f16.f.v2bf16(<2 x bfloat> %_)
|
||||
static METAL_FUNC half2 convert(device const bf16_2 & u) { return half2(u); }
|
||||
static METAL_FUNC half2 convert(threadgroup const bf16_2 & u) { return half2(u); }
|
||||
static METAL_FUNC half2 convert(thread const bf16_2 & u) { return half2(u); }
|
||||
};
|
||||
|
||||
|
||||
|
||||
} // base_types
|
||||
|
||||
} // mittens
|
||||
@@ -0,0 +1,10 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief A collection of common resources on which Thundermittens depends.
|
||||
*/
|
||||
|
||||
|
||||
#pragma once
|
||||
#include "base_types.metal"
|
||||
#include "base_ops.metal"
|
||||
#include "utils.metal"
|
||||
@@ -0,0 +1,225 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief General utilities for Thundermittens.
|
||||
*/
|
||||
#pragma once // not done
|
||||
/*
|
||||
TODO:
|
||||
shared allocator
|
||||
max shared mem for other hardware
|
||||
*/
|
||||
|
||||
#include <metal_stdlib>
|
||||
#include "base_types.metal"
|
||||
/**
|
||||
* @namespace mittens
|
||||
*
|
||||
* @brief The main namespace of Thundermittens.
|
||||
*/
|
||||
namespace mittens {
|
||||
/**
|
||||
* @namespace ore
|
||||
*
|
||||
* @brief The main namespace of Thundermittens Metal.
|
||||
*/
|
||||
|
||||
/* ---------- GENERAL CONSTANTS FOR mittens ---------- */
|
||||
|
||||
/**
|
||||
* @brief Tile dimension constant.
|
||||
*/
|
||||
constant constexpr const int TILE_DIM{8};
|
||||
constant constexpr const int TILE_ELEMENTS{TILE_DIM*TILE_DIM};
|
||||
constant constexpr const int SIMD_THREADS{32};
|
||||
|
||||
|
||||
#ifdef M2_PRO
|
||||
constant constexpr int MAX_SHARED_MEMORY = 32768;
|
||||
#else
|
||||
constant constexpr int MAX_SHARED_MEMORY = 32768;
|
||||
#endif
|
||||
/* ---------- TYPE HELPERS ---------- */
|
||||
/**
|
||||
* @namespace ducks
|
||||
*
|
||||
* @brief Thundermittens' 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 mittens.
|
||||
#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.
|
||||
*/
|
||||
constant static constexpr uint32_t MASK_ALL = 0xFFFFFFFF;
|
||||
|
||||
template<typename T>
|
||||
static METAL_FUNC T shfl_sync(thread const T &f, const ushort laneid) {
|
||||
return metal::simd_shuffle(f, laneid);
|
||||
}
|
||||
|
||||
template<>
|
||||
METAL_FUNC bfloat shfl_sync<bfloat>(thread const bf16 &f, const ushort laneid) {
|
||||
// return as_type<bf16>(metal::simd_shuffle(*(thread half*)(&f), laneid));
|
||||
float f_val = (float)f;
|
||||
float shfl_val = metal::simd_shuffle(f_val, laneid);
|
||||
return (bf16)shfl_val;
|
||||
}
|
||||
|
||||
template<>
|
||||
METAL_FUNC bfloat2 shfl_sync<bfloat2>(thread const bf16_2 &f, const ushort laneid) {
|
||||
// return as_type<bf16_2>(metal::simd_shuffle(*(thread half2*)(&f), laneid));
|
||||
float2 f_val = (float2)f;
|
||||
float2 shfl_val = metal::simd_shuffle(f_val, laneid);
|
||||
return (bf16_2)shfl_val;
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
static METAL_FUNC T shfl_down_fill_sync(thread const T &f, thread const T& fill_data, const ushort laneid) {
|
||||
return metal::simd_shuffle_and_fill_down(f, laneid, fill_data);
|
||||
}
|
||||
|
||||
template<>
|
||||
METAL_FUNC bfloat shfl_down_fill_sync<bfloat>(thread const bfloat &f, thread const bfloat &fill_data, const ushort laneid) {
|
||||
// return as_type<bf16>(metal::simd_shuffle_and_fill_down(*(thread half*)(&f), *(thread half*)(&fill_data), laneid));
|
||||
float f_val = (float)f;
|
||||
float fill_data_f = (float)fill_data;
|
||||
float shfl_val = metal::simd_shuffle_and_fill_down(f_val, fill_data_f, laneid);
|
||||
return (bf16)shfl_val;
|
||||
}
|
||||
template<>
|
||||
METAL_FUNC bfloat2 shfl_down_fill_sync<bfloat2>(thread const bfloat2 &f, thread const bfloat2 &fill_data, const ushort laneid) {
|
||||
// return as_type<bf16_2>(metal::simd_shuffle_and_fill_down(*(thread half2*)(&f), *(thread half2*)(&fill_data), laneid));
|
||||
float2 f_val = (float2)f;
|
||||
float2 fill_data_f = (float2)fill_data;
|
||||
float2 shfl_val = metal::simd_shuffle_and_fill_down(f_val, fill_data_f, laneid);
|
||||
return (bf16_2)shfl_val;
|
||||
}
|
||||
/**
|
||||
* @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>
|
||||
static METAL_FUNC T shfl_down_sync(thread const T &f, int delta) {
|
||||
return metal::simd_shuffle_rotate_down(f, delta);
|
||||
}
|
||||
|
||||
template<>
|
||||
METAL_FUNC bfloat shfl_down_sync<bfloat>(thread const bf16 &f, int delta) {
|
||||
// return base_types::convertor<bf16, float>::convert(metal::simd_shuffle_rotate_down(base_types::convertor<float, bf16>::convert(f), delta));
|
||||
// return as_type<bf16>(metal::simd_shuffle_rotate_down(*(thread half*)(&f), delta));
|
||||
float f_val = (float)f;
|
||||
float shfl_val = metal::simd_shuffle_rotate_down(f_val, delta);
|
||||
return (bf16)shfl_val;
|
||||
}
|
||||
|
||||
template<>
|
||||
METAL_FUNC bfloat2 shfl_down_sync<bfloat2>(thread const bf16_2 &f, int delta) {
|
||||
// return as_type<bf16_2>(metal::simd_shuffle_rotate_down(*(thread const half2*)(&f), delta));
|
||||
// return base_types::convertor<bf16_2, float2>::convert(metal::simd_shuffle_rotate_down(base_types::convertor<float2, bf16_2>::convert(f), delta));
|
||||
|
||||
float2 f_val = (float2)f;
|
||||
float2 shfl_val = metal::simd_shuffle_rotate_down(f_val, delta);
|
||||
return (bf16_2)shfl_val;
|
||||
// return as_type<bf16_2>(metal::simd_shuffle_rotate_down(*(thread half2*)(&f), delta));
|
||||
}
|
||||
|
||||
|
||||
/* ---------- LOOP UNROLLING UTILS ---------- */
|
||||
|
||||
namespace meta {
|
||||
template <int Start, int End, int Stride, bool = (Start < End)>
|
||||
struct unroll_i_in_range {
|
||||
template<class F, typename... Args>
|
||||
static METAL_FUNC void run(F f, Args... args) {
|
||||
f(Start, args...);
|
||||
unroll_i_in_range<Start + Stride, End, Stride>::run(f, args...);
|
||||
}
|
||||
};
|
||||
|
||||
template <int Start, int End, int Stride>
|
||||
struct unroll_i_in_range<Start, End, Stride, false> {
|
||||
template<class F, typename... Args>
|
||||
static METAL_FUNC void run(F, Args...) {
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
template <int Start, int End, int Stride, bool = (Start < End)>
|
||||
struct unroll_i_j_in_range_inner {
|
||||
template<class F, typename... Args>
|
||||
static METAL_FUNC void run(F f, int outerIndex, Args... args) {
|
||||
f(outerIndex, Start, args...);
|
||||
unroll_i_j_in_range_inner<Start + Stride, End, Stride>::run(f, outerIndex, args...);
|
||||
}
|
||||
};
|
||||
|
||||
template <int Start, int End, int Stride>
|
||||
struct unroll_i_j_in_range_inner<Start, End, Stride, false> {
|
||||
template<class F, typename... Args>
|
||||
static METAL_FUNC void run(F, int, Args...) {
|
||||
}
|
||||
};
|
||||
|
||||
template <int StartOuter, int EndOuter, int StrideOuter,
|
||||
int StartInner, int EndInner, int StrideInner,
|
||||
bool = (StartOuter < EndOuter)>
|
||||
struct unroll_i_j_in_range {
|
||||
template<class F, typename... Args>
|
||||
static METAL_FUNC void run(F f, Args... args) {
|
||||
unroll_i_j_in_range_inner<StartInner, EndInner, StrideInner>::run(
|
||||
f, StartOuter, args...
|
||||
);
|
||||
unroll_i_j_in_range<
|
||||
StartOuter + StrideOuter, EndOuter, StrideOuter,
|
||||
StartInner, EndInner, StrideInner
|
||||
>::run(f, args...);
|
||||
}
|
||||
};
|
||||
|
||||
template <int StartOuter, int EndOuter, int StrideOuter,
|
||||
int StartInner, int EndInner, int StrideInner>
|
||||
struct unroll_i_j_in_range<StartOuter, EndOuter, StrideOuter,
|
||||
StartInner, EndInner, StrideInner, false> {
|
||||
template<class F, typename... Args>
|
||||
static METAL_FUNC void run(F, Args...) {
|
||||
}
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
|
||||
template <int N>
|
||||
struct ReadVector {
|
||||
float _[N];
|
||||
};
|
||||
|
||||
/* ---------- SHARED MEMORY UTILS ---------- */
|
||||
|
||||
#define mittens_ALIGN_AS(n) alignas(n)
|
||||
#define mittens_DEFAULT_ALIGN mittens_ALIGN_AS(16)
|
||||
|
||||
/**
|
||||
* @brief Dummy structure for alignment purposes. Needed for WGMMA and TMA calls.
|
||||
*/
|
||||
struct mittens_DEFAULT_ALIGN alignment_dummy { int dummy; };
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief An aggregate header of all group (multi-warp) operations defined by Thundermittens
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
#include "../../common/common.metal"
|
||||
#include "../../types/types.metal"
|
||||
#include "../warp/warp.metal" // several group memory ops rely on underlying warp-scope ops
|
||||
namespace mittens {
|
||||
template<int N_WARPS>
|
||||
struct group {
|
||||
constant static constexpr int GROUP_WARPS = N_WARPS; // This alias produces nice parallelism.
|
||||
constant static constexpr int GROUP_THREADS = N_WARPS * mittens::SIMD_THREADS; // This alias produces nice parallelism.
|
||||
static METAL_FUNC int simd_laneid(const unsigned threadIdx) { return threadIdx % mittens::SIMD_THREADS; }
|
||||
static METAL_FUNC int laneid (const unsigned threadIdx) { return threadIdx % GROUP_THREADS; }
|
||||
static METAL_FUNC int warpid (const unsigned threadIdx) { return laneid(threadIdx) / mittens::SIMD_THREADS; }
|
||||
static METAL_FUNC int groupid (const unsigned threadIdx) { return threadIdx / GROUP_THREADS; }
|
||||
#include "memory/memory.metal"
|
||||
#include "shared/shared.metal"
|
||||
};
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
#include "tile/tile.metal"
|
||||
#include "vec/vec.metal"
|
||||
@@ -0,0 +1,132 @@
|
||||
|
||||
/**
|
||||
* @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<typename RT, typename GL>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
|
||||
load(thread RT &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename GL::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
const device U *src = (device U*)&_src.template get<RT>(idx);
|
||||
const int row_stride = _src.row_stride();
|
||||
|
||||
int warp_laneid = threadIdx % 32;
|
||||
const int row_offset = dst.rows * warpid(threadIdx);
|
||||
const short qid = warp_laneid / 4;
|
||||
const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
|
||||
const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
int row = simd_y + i * RT::tile_size;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
int col = simd_x + j * RT::tile_size;
|
||||
T2 src2 = base_types::convertor<T2, U2>::convert(*((device U2*)(&src[row * row_stride + col])));
|
||||
dst.tiles[i][j].data.thread_elements()[0] = src2[0];
|
||||
dst.tiles[i][j].data.thread_elements()[1] = src2[1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<typename RT, typename GL>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
|
||||
load(thread RT &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename GL::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
const device U *src = (device U*)&_src.template get<RT>(idx);
|
||||
const int row_stride = _src.row_stride();
|
||||
|
||||
int warp_laneid = threadIdx % 32;
|
||||
const int row_offset = dst.rows * warpid(threadIdx);
|
||||
const short qid = warp_laneid / 4;
|
||||
const short simd_y = row_offset + (qid & 2) * 2 + (warp_laneid % 2) * 2;;
|
||||
const short simd_x = (qid & 4) + (warp_laneid / 2) % 4;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
int row = simd_y + i * RT::tile_size;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
int col = simd_x + j * RT::tile_size;
|
||||
T2 src2 = base_types::convertor<T2, U2>::convert(*((device U2*)(&src[row * row_stride + col])));
|
||||
dst.tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src[row * row_stride + col]);
|
||||
dst.tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src[(row + 1) * row_stride + col]);
|
||||
}
|
||||
}
|
||||
}
|
||||
/**
|
||||
* @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<typename RT, typename GL>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>(), void>::type
|
||||
store(thread GL &_dst, thread const RT &src, thread const coord &idx, const int threadIdx) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename GL::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
device U *dst = (device U*)&(_dst.template get<RT>(idx));
|
||||
const int row_stride = _dst.row_stride();
|
||||
int warp_laneid = simd_laneid(threadIdx);
|
||||
const int row_offset = src.rows * warpid(threadIdx);
|
||||
const short qid = warp_laneid / 4;
|
||||
const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
|
||||
const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < src.height; i++) {
|
||||
int row = simd_y + i * RT::tile_size;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < src.width; j++) {
|
||||
int col = simd_x + j * RT::tile_size;
|
||||
U2 src2 = base_types::convertor<U2, T2>::convert(T2(src.tiles[i][j].data.thread_elements()[0], src.tiles[i][j].data.thread_elements()[1]));
|
||||
*(device U2*)(&dst[row*row_stride + col]) = src2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<typename RT, typename GL>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_col_register_tile<RT>(), void>::type
|
||||
store(thread GL &_dst, thread const RT &src, thread const coord &idx, const int threadIdx) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename GL::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
device U *dst = (device U*)&(_dst.template get<RT>(idx));
|
||||
const int row_stride = _dst.row_stride();
|
||||
int warp_laneid = simd_laneid(threadIdx);
|
||||
const int row_offset = src.rows * warpid(threadIdx);
|
||||
const short qid = warp_laneid / 4;
|
||||
// const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
|
||||
// const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
|
||||
const short simd_y = row_offset + (qid & 2) * 2 + (warp_laneid % 2) * 2;
|
||||
const short simd_x = (qid & 4) + (warp_laneid / 2) % 4;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < src.height; i++) {
|
||||
int row = simd_y + i * RT::tile_size;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < src.width; j++) {
|
||||
int col = simd_x + j * RT::tile_size;
|
||||
dst[row*row_stride + col] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[0]);
|
||||
dst[(row + 1) * row_stride + col] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[1]);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Group (collaborative warp) ops for loading shared tiles from and storing to global memory.
|
||||
*/
|
||||
|
||||
|
||||
//template<typename ST, typename U>
|
||||
//static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
//load(int i,
|
||||
// threadgroup ST *dst, device U* src,
|
||||
// thread const int& group_laneid,
|
||||
// thread const int& memcpy_per_row,
|
||||
// thread const int& elem_per_memcpy,
|
||||
// thread const int& row_stride)
|
||||
//{
|
||||
// int idx = i * GROUP_THREADS + group_laneid;
|
||||
// int row = idx / memcpy_per_row;
|
||||
// int col = (idx*elem_per_memcpy) % ST::cols;
|
||||
// if (row < ST::rows) {
|
||||
// *(threadgroup float4*)(&(*dst)[{row, col}]) = *(device float4*)(&src[row*row_stride + col]);
|
||||
// }
|
||||
//}
|
||||
|
||||
|
||||
template<typename ST, typename GL>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
|
||||
load(threadgroup ST &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
|
||||
int group_laneid = threadIdx % GROUP_THREADS;
|
||||
using T = typename ST::T;
|
||||
using U = typename GL::dtype;
|
||||
device U *src = (device U*)&_src.template get<ST>(idx);
|
||||
const int row_stride = _src.row_stride();
|
||||
using read_vector = ReadVector<1>;
|
||||
// we can handle this many rows each time we run a memcpy_async
|
||||
constexpr const int elem_per_memcpy = sizeof(read_vector)/sizeof(typename ST::dtype);
|
||||
constexpr const int memcpy_per_row = ST::cols / elem_per_memcpy;
|
||||
int total_calls = ((ST::height * ST::width + (N_WARPS-1))) * TILE_DIM*TILE_DIM / (N_WARPS*SIMD_THREADS*elem_per_memcpy); // round up
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < total_calls; i++) {
|
||||
|
||||
int idx = i * GROUP_THREADS + group_laneid;
|
||||
int row = idx / memcpy_per_row;
|
||||
int col = (idx*elem_per_memcpy) % dst.cols;
|
||||
if (row<dst.rows && col < dst.cols) {
|
||||
*(threadgroup read_vector*)(&dst[{row, col}]) = *(device read_vector*)(&src[row*row_stride + col]);
|
||||
// *(threadgroup float*)(&dst[{row, col}]) = 1.0f;
|
||||
}
|
||||
}
|
||||
// dst[{0, 0}] = base_types::convertor<T, float>::convert(1.f);
|
||||
// dst[{0, 0}] = total_calls;
|
||||
// meta::unroll_i_in_range<0, total_calls, 1>::run(load<ST, typename GL::dtype>, &dst, src, group_laneid, memcpy_per_row, elem_per_memcpy, row_stride);
|
||||
}
|
||||
|
||||
|
||||
//template<typename ST, typename GL>
|
||||
//static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
|
||||
//load(threadgroup ST &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
|
||||
// int group_laneid = threadIdx % GROUP_THREADS;
|
||||
// int groupid = threadIdx / GROUP_THREADS;
|
||||
// int laneid = threadIdx % SIMD_THREADS;
|
||||
//
|
||||
// using U = typename GL::dtype;
|
||||
// device U *src = (device U*)&_src.template get<ST>(idx);
|
||||
// const int row_stride = _src.row_stride();
|
||||
//
|
||||
// int elem_per_memcpy = sizeof(float)/sizeof(typename ST::dtype);
|
||||
// int memcpy_per_row = ST::cols / elem_per_memcpy;
|
||||
// int total_calls = ((ST::height * ST::width + (N_WARPS-1))) * TILE_DIM*TILE_DIM / (N_WARPS*SIMD_THREADS*elem_per_memcpy); // round up
|
||||
// /*
|
||||
// 1x16 or 8 x 128
|
||||
// */
|
||||
// int offset = ST::num_elements / (GROUP_WARPS);
|
||||
//// int offset = group_laneid
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < total_calls; i++) {
|
||||
// int idx = i * SIMD_THREADS + laneid;
|
||||
//// int idx = i * () + group_laneid;
|
||||
// int row = idx / memcpy_per_row;
|
||||
// int col = (idx*elem_per_memcpy) % dst.cols;
|
||||
// if (row<dst.rows) {
|
||||
// *(threadgroup float*)(&dst[{row, col}]) = *(device float*)(&src[row*row_stride + col]);
|
||||
// }
|
||||
// }
|
||||
//}
|
||||
//
|
||||
//template<typename ST, typename GL>
|
||||
//static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
|
||||
//load(threadgroup ST &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
|
||||
// int warp_id = threadIdx / SIMD_THREADS;
|
||||
// int lane_id = threadIdx % SIMD_THREADS;
|
||||
//// int N_WARPS = /* number of warps in your group */;
|
||||
//
|
||||
// using U = typename GL::dtype;
|
||||
// device U *src = (device U*)&_src.template get<ST>(idx);
|
||||
// const int row_stride = _src.row_stride();
|
||||
//
|
||||
// int elem_per_memcpy = sizeof(float)/sizeof(typename ST::dtype);
|
||||
// int memcpy_per_row = ST::cols / elem_per_memcpy;
|
||||
// int total_memcpy_elems = (ST::height * ST::cols) / elem_per_memcpy;
|
||||
// int elems_per_warp = (total_memcpy_elems + N_WARPS - 1) / N_WARPS; // Ceiling division
|
||||
//
|
||||
// int start_idx = warp_id * elems_per_warp;
|
||||
// int end_idx = metal::min(start_idx + elems_per_warp, total_memcpy_elems);
|
||||
//
|
||||
// #pragma clang loop unroll(full)
|
||||
// for (int idx = start_idx + lane_id; idx < end_idx; idx += SIMD_THREADS) {
|
||||
// int row = idx / memcpy_per_row;
|
||||
// int col = (idx % memcpy_per_row) * elem_per_memcpy;
|
||||
// if (row < ST::height) {
|
||||
// *(threadgroup float*)(&dst[{row, col}]) = *(device float*)(&src[row * row_stride + col]);
|
||||
// }
|
||||
// }
|
||||
//}
|
||||
|
||||
template<typename ST, typename GL>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
|
||||
store(thread const GL &_dst, threadgroup const ST &src, thread const coord &idx, const int threadIdx) {
|
||||
int group_laneid = threadIdx % GROUP_THREADS;
|
||||
using U = typename GL::dtype;
|
||||
device U *dst = (device U*)&_dst.template get<ST>(idx);
|
||||
const int row_stride = _dst.row_stride();
|
||||
using read_vector = ReadVector<1>;
|
||||
// we can handle this many rows each time we run a memcpy_async
|
||||
int elem_per_memcpy = sizeof(read_vector)/sizeof(typename ST::dtype); // float/float -> 1
|
||||
int memcpy_per_row = ST::cols / elem_per_memcpy; // 240 memcpy per row
|
||||
int total_calls = ((src.height * src.width + (N_WARPS-1))) * TILE_DIM*TILE_DIM / (N_WARPS*SIMD_THREADS*elem_per_memcpy); // round up
|
||||
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < total_calls; i++) {
|
||||
|
||||
int idx = i * GROUP_THREADS + group_laneid;
|
||||
|
||||
int row = idx / memcpy_per_row;
|
||||
int col = (idx*elem_per_memcpy) % src.cols;
|
||||
if (row<src.rows && col < src.cols) {
|
||||
*(device read_vector*)(&dst[row*row_stride + col]) = *(threadgroup read_vector*)(&src[{row, col}]);
|
||||
// *(device float*)(&dst[row*row_stride + col]) = 1.f;
|
||||
}
|
||||
}
|
||||
// dst[0] = src[{0,0}];
|
||||
// dst[0] = total_calls;
|
||||
// dst[0] = base_types::convertor<U, float>::convert(1);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
/**
|
||||
* @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<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
load(thread RT &dst, threadgroup const ST &src, const int threadIdx) {
|
||||
constexpr int height = ST::height;
|
||||
constexpr int warp_height = RT::height;
|
||||
static_assert(height%N_WARPS == 0, "Group load / store requires tile height to be a multiple of N_WARPS.");
|
||||
static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
|
||||
static_assert(warp_height * N_WARPS == height, "RT height * N_WARPS must = ST height");
|
||||
static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
|
||||
int warp_laneid = simd_laneid(threadIdx);
|
||||
const int row_offset = RT::rows * warpid(threadIdx);
|
||||
const short qid = warp_laneid / 4;
|
||||
const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
|
||||
const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
int row = simd_y + i * mittens::TILE_DIM;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
int col = simd_x + j * mittens::TILE_DIM;
|
||||
T2 src2 = base_types::convertor<T2, U2>::convert(*((threadgroup U2*)(&src[{row, col}])));
|
||||
dst.tiles[i][j].data.thread_elements()[0] = src2[0];
|
||||
dst.tiles[i][j].data.thread_elements()[1] = src2[1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
load(thread RT &dst, threadgroup const ST &src, const int threadIdx) {
|
||||
constexpr int height = ST::height;
|
||||
constexpr int warp_height = RT::height;
|
||||
static_assert(height%N_WARPS == 0, "Group load / store requires tile height to be a multiple of N_WARPS.");
|
||||
static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
|
||||
static_assert(warp_height * N_WARPS == height, "RT height * N_WARPS must = ST height");
|
||||
static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
|
||||
int warp_laneid = simd_laneid(threadIdx);
|
||||
const int row_offset = RT::rows * warpid(threadIdx);
|
||||
const short qid = warp_laneid / 4;
|
||||
const short simd_y = row_offset + (qid & 2) * 2 + (warp_laneid % 2) * 2;
|
||||
const short simd_x = (qid & 4) + (warp_laneid / 2) % 4;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
int row = simd_y + i * mittens::TILE_DIM;
|
||||
int col = simd_x + j * mittens::TILE_DIM;
|
||||
dst.tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src[{row + 0, col}]);
|
||||
dst.tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src[{row + 1, col}]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @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<typename ST, typename RT>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
store(threadgroup ST &dst, thread const RT &src, const int threadIdx) {
|
||||
constexpr int height = ST::height;
|
||||
constexpr int warp_height = RT::height;
|
||||
static_assert(height%N_WARPS == 0, "Group load / store requires tile height to be a multiple of N_WARPS.");
|
||||
static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
|
||||
static_assert(warp_height * N_WARPS == height, "RT height * N_WARPS must = ST height");
|
||||
static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
int warp_laneid = simd_laneid(threadIdx);
|
||||
const int row_offset = RT::rows * warpid(threadIdx);
|
||||
const short qid = warp_laneid / 4;
|
||||
const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
|
||||
const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < RT::height; i++) {
|
||||
int row = simd_y + i * mittens::TILE_DIM;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < RT::width; j++) {
|
||||
int col = simd_x + j * mittens::TILE_DIM;
|
||||
U2 src2 = base_types::convertor<U2, T2>::convert(T2(src.tiles[i][j].data.thread_elements()[0],
|
||||
src.tiles[i][j].data.thread_elements()[1]));
|
||||
*(threadgroup U2*)(&dst[{row, col}]) = src2;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template<typename ST, typename RT>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
store(threadgroup ST &dst, thread const RT &src, const int threadIdx) {
|
||||
constexpr int height = ST::height;
|
||||
constexpr int warp_height = RT::height;
|
||||
static_assert(height%N_WARPS == 0, "Group load / store requires tile height to be a multiple of N_WARPS.");
|
||||
static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
|
||||
static_assert(warp_height * N_WARPS == height, "RT height * N_WARPS must = ST height");
|
||||
static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
int warp_laneid = simd_laneid(threadIdx);
|
||||
const int row_offset = RT::rows * warpid(threadIdx);
|
||||
const short qid = warp_laneid / 4;
|
||||
// const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
|
||||
// const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
|
||||
const short simd_y = row_offset + (qid & 2) * 2 + (warp_laneid % 2) * 2;
|
||||
const short simd_x = (qid & 4) + (warp_laneid / 2) % 4;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < RT::height; i++) {
|
||||
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < RT::width; j++) {
|
||||
int row = simd_y + i * mittens::TILE_DIM;
|
||||
int col = simd_x + j * mittens::TILE_DIM;
|
||||
// U2 src2 = base_types::convertor<U2, T2>::convert(T2(src.tiles[i][j].data.thread_elements()[0],
|
||||
// src.tiles[i][j].data.thread_elements()[1]));
|
||||
// *(threadgroup U2*)(&dst[{row, col}]) = src2;
|
||||
|
||||
dst[{row + 0, col}] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[0]);
|
||||
dst[{row + 1, col}] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[1]);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief An aggregate header of group memory operations on tiles.
|
||||
*/
|
||||
|
||||
#include "shared_to_register.metal"
|
||||
#include "global_to_register.metal"
|
||||
#include "global_to_shared.metal"
|
||||
@@ -0,0 +1,47 @@
|
||||
|
||||
/**
|
||||
* @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<typename RV, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
load(thread RV &dst, thread const GL &_src, thread coord idx, const int threadIdx) {
|
||||
using T = typename RV::dtype;
|
||||
using U = typename GL::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
|
||||
idx.c += warpid(threadIdx);
|
||||
// Call warp level store
|
||||
::mittens::load(dst, _src, idx, simd_laneid(threadIdx));
|
||||
}
|
||||
|
||||
/**
|
||||
* @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<typename RV, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
store(thread GL &_dst, thread const RV &src, thread coord idx, const int threadIdx) {
|
||||
using T = typename RV::dtype;
|
||||
// using U2 = typename base_types::packing<U>::packed_type;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
|
||||
idx.c += warpid(threadIdx);
|
||||
|
||||
// Call warp level store
|
||||
::mittens::store(_dst, src, idx, simd_laneid(threadIdx));
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
/**
|
||||
* @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<typename SV, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
load(threadgroup SV &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
|
||||
using U = typename GL::dtype;
|
||||
using read_vector = ReadVector<1>;
|
||||
constexpr int elem_per_transfer = sizeof(read_vector) / sizeof(typename SV::dtype);
|
||||
constexpr int total_calls = SV::length / elem_per_transfer; // guaranteed to divide
|
||||
device U *src = (device U*)&_src.template get<SV>(idx);
|
||||
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid(threadIdx); i < total_calls; i+=GROUP_THREADS) {
|
||||
if(i * elem_per_transfer < dst.length)
|
||||
*(threadgroup read_vector*)&dst[i*elem_per_transfer] = *(device read_vector*)&src[i*elem_per_transfer];
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @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<typename SV, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
store(thread const GL &_dst, threadgroup const SV &src, thread const coord &idx, const int threadIdx) {
|
||||
using read_vector = ReadVector<1>;
|
||||
using U = typename GL::dtype;
|
||||
constexpr int elem_per_transfer = sizeof(read_vector) / sizeof(typename SV::dtype);
|
||||
constexpr int total_calls = SV::length / elem_per_transfer; // guaranteed to divide
|
||||
device U *dst = (device U*)&_dst.template get<SV>(idx);
|
||||
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid(threadIdx); i < total_calls; i+= GROUP_THREADS) {
|
||||
if(i * elem_per_transfer < src.length)
|
||||
*(device read_vector*)&dst[i*elem_per_transfer] = *(threadgroup read_vector*)&src[i*elem_per_transfer]; // lmao it's identical
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
/**
|
||||
* @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<typename RV, typename SV>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
load(thread RV &dst, threadgroup const SV &_src, const int threadIdx) {
|
||||
using T = typename RV::dtype;
|
||||
using U = typename SV::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
|
||||
static_assert(SV::length == RV::length*N_WARPS, "rv and sv dimensions do not match");// confirm size correct
|
||||
// threadgroup typename SV::template subvec<typename SV::dtype, RV::outer_dim> &src = subvec_inplace<RV::outer_dim, SV>(_src, warpid(threadIdx));
|
||||
// threadgroup subvec &src = subvec_inplace<RV::outer_dim, SV>(_src, warpid(threadIdx));
|
||||
unsigned warpId = warpid(threadIdx);
|
||||
using subvec = typename SV::template subvec<RV::length>;
|
||||
|
||||
threadgroup subvec& src = *(threadgroup subvec*)(&_src[warpId *RV::length]);
|
||||
|
||||
::mittens::load<RV, subvec>(dst, src, simd_laneid(threadIdx)); // 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<typename SV, typename RV>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
store(threadgroup SV &_dst, thread const RV &src, const int threadIdx) {
|
||||
using T = typename RV::dtype;
|
||||
using U = typename SV::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
|
||||
|
||||
static_assert(SV::length == RV::length*N_WARPS, "rv and sv dimensions do not match");// confirm size correct
|
||||
|
||||
// threadgroup typename SV::template subvec<typename SV::dtype, RV::outer_dim> &dst = subvec_inplace<RV::outer_dim, SV>(_dst, warpid(threadIdx));
|
||||
// ::mittens::store<threadgroup typename SV::template subvec<typename SV::dtype, RV::outer_dim>, RV>(dst, src, simd_laneid(threadIdx)); // warp-level
|
||||
|
||||
unsigned warpId = warpid(threadIdx);
|
||||
using subvec = typename SV::template subvec<RV::length>;
|
||||
threadgroup subvec& dst = *(threadgroup subvec*)(&_dst[warpId * RV::length]);
|
||||
|
||||
::mittens::store(dst, src, simd_laneid(threadIdx)); // warp-level
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief An aggregate header of group memory operations on vectors.
|
||||
*/
|
||||
|
||||
#include "shared_to_register.metal"
|
||||
#include "global_to_register.metal"
|
||||
#include "global_to_shared.metal"
|
||||
@@ -0,0 +1,3 @@
|
||||
|
||||
#include "tile/tile.metal"
|
||||
#include "vec/vec.metal"
|
||||
@@ -0,0 +1,27 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Group conversions between different shared memory tile types.
|
||||
*/
|
||||
|
||||
/* ---------- COPIES ---------- */
|
||||
|
||||
/**
|
||||
* @brief Copies data from one shared memory tile to another, potentially with different data types and layouts.
|
||||
*
|
||||
* @tparam T The data type of the destination tile.
|
||||
* @tparam U The data type of the source tile.
|
||||
* @tparam _height The height of the tile.
|
||||
* @tparam _width The width of the tile.
|
||||
* @tparam L1 The layout of the destination tile.
|
||||
* @tparam L2 The layout of the source tile.
|
||||
* @param[out] dst The destination tile.
|
||||
* @param[in] src The source tile.
|
||||
*/
|
||||
template<typename T, typename U, int _height, int _width>
|
||||
static METAL_FUNC void copy(threadgroup st<T, _height, _width> &dst, threadgroup const st<U, _height, _width> &src, const int threadIdx) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid(threadIdx); i < dst.num_elements; i+=GROUP_THREADS) {
|
||||
int row = i/dst.cols, col = i%dst.cols;
|
||||
dst[{row, col}] = base_types::convertor<T, U>::convert(src[{row, col}]);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,475 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Group maps on shared tiles.
|
||||
*/
|
||||
|
||||
/**
|
||||
* @brief Performs a uniform unary operation on a tile.
|
||||
*
|
||||
* This function applies a given unary operation to each element of the source tile and stores the result in the destination tile.
|
||||
* The operation is applied independently to each element, without considering its position or the values of neighboring elements.
|
||||
*
|
||||
* @tparam op The unary operation to be applied. Must be specialized to support operation on the data type of T.
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the unary operation is applied.
|
||||
*/
|
||||
template<typename op, typename ST> // T2, w, h can be inferred from dst as long as op is specialized
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
unary_map(threadgroup ST &dst, threadgroup const ST &src, const int threadIdx) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_THREADS) {
|
||||
dst.data[i] = op::template op<typename ST::dtype>(src.data[i]);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Performs a uniform binary operation on a tile with a scalar parameter.
|
||||
*
|
||||
* This function applies a given binary operation to each element of the source tile and a scalar parameter, then stores the result in the destination tile.
|
||||
* The operation is applied independently to each element, treating the scalar parameter as the second operand for each operation.
|
||||
*
|
||||
* @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T and the scalar parameter.
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the binary operation is applied.
|
||||
* @param[in] param The scalar parameter to be used as the second operand in the binary operation.
|
||||
*/
|
||||
template<typename op, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
bin_map(threadgroup ST &dst, threadgroup const ST &src, thread const typename ST::dtype ¶m, const int threadIdx) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_THREADS) {
|
||||
dst.data[i] = op::template op<typename ST::dtype>(src.data[i], param);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Performs a uniform binary operation on two tiles.
|
||||
*
|
||||
* This function applies a given binary operation to corresponding elements of two source tiles and stores the result in the destination tile.
|
||||
* The operation is applied independently to each pair of elements, without considering their positions or the values of neighboring elements.
|
||||
*
|
||||
* @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T.
|
||||
* @tparam T The type of the tiles. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile to which the binary operation is applied.
|
||||
* @param[in] rhs The second source tile to which the binary operation is applied.
|
||||
*/
|
||||
template<typename op, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
bin_map(threadgroup ST &dst, threadgroup const ST &lhs, threadgroup const ST &rhs, const int threadIdx) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_THREADS) {
|
||||
dst.data[i] = op::template op<typename ST::dtype>(lhs.data[i], rhs.data[i]);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Performs a row-wise binary operation on a tile with a vector.
|
||||
*
|
||||
* This function applies a given binary operation to each row of the source tile and the corresponding element of the source vector,
|
||||
* then stores the result in the destination tile. The operation is applied independently to each row, using the vector element as
|
||||
* the second operand for each element in the row.
|
||||
*
|
||||
* @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T and the vector elements.
|
||||
* @tparam T The type of the tiles. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam V The type of the vector. Must have the same data type as T.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the binary operation is applied.
|
||||
* @param[in] vec The source vector containing the second operand for each row operation.
|
||||
*/
|
||||
template<typename op, typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_map(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &vec, const int threadIdx) {
|
||||
static_assert(metal::is_same<typename ST::dtype, typename SV::dtype>::value, "Tile and vector must have the same data type");
|
||||
static_assert(SV::length == ST::rows, "Vector length must match the number of rows in the tile");
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_THREADS) {
|
||||
int row = i/dst.cols, col = i%dst.cols;
|
||||
dst[{row, col}] = op::template op<typename ST::dtype>(src[{row, col}], vec[row]);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Performs a column-wise binary operation on a tile with a vector.
|
||||
*
|
||||
* This function applies a given binary operation to each column of the source tile and the corresponding element of the source vector,
|
||||
* then stores the result in the destination tile. The operation is applied independently to each column, using the vector element as
|
||||
* the second operand for each element in the column.
|
||||
*
|
||||
* @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T and the vector elements.
|
||||
* @tparam T The type of the tiles. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam V The type of the vector. Must have the same data type as T.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the binary operation is applied.
|
||||
* @param[in] vec The source vector containing the second operand for each column operation.
|
||||
*/
|
||||
template<typename op, typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_map(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &vec, const int threadIdx) {
|
||||
static_assert(metal::is_same<typename ST::dtype, typename SV::dtype>::value, "Tile and vector must have the same data type");
|
||||
static_assert(SV::length == ST::cols, "Vector length must match the number of columns in the tile");
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_THREADS) {
|
||||
int row = i/dst.cols, col = i%dst.cols;
|
||||
dst[{row, col}] = op::template op<typename ST::dtype>(src[{row, col}], vec[col]);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
|
||||
|
||||
// All of the annoying qualifiers *should* be automatically inferred during compile-time.
|
||||
// So, syntax should just be mittens::add_row(tile, colvec);
|
||||
|
||||
// const maps
|
||||
/**
|
||||
* @brief Sets all elements of the destination tile to zero.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
zero(threadgroup ST &dst, const int threadIdx) {
|
||||
unary_map<base_ops::zero, ST>(dst, dst, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of the destination tile to one.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
one(threadgroup ST &dst, const int threadIdx) {
|
||||
unary_map<base_ops::one, ST>(dst, dst, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of the destination tile to positive infinity.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
pos_infty(threadgroup ST &dst, const int threadIdx) {
|
||||
unary_map<base_ops::pos_infty, ST>(dst, dst, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of the destination tile to negative infinity.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
neg_infty(threadgroup ST &dst, const int threadIdx) {
|
||||
unary_map<base_ops::neg_infty, ST>(dst, dst, threadIdx);
|
||||
}
|
||||
|
||||
// unary maps
|
||||
/**
|
||||
* @brief Applies the exponential function to each element of the source tile and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the exponential function is applied.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
exp(threadgroup ST &dst, threadgroup const ST &src, const int threadIdx) {
|
||||
unary_map<base_ops::exp, ST>(dst, src, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the exponential function to each element of the source tile and stores the result in the destination tile, in base 2.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the exponential function is applied.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
exp2(threadgroup ST &dst, threadgroup const ST &src, const int threadIdx) {
|
||||
unary_map<base_ops::exp2, ST>(dst, src, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the natural logarithm function to each element of the source tile and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the natural logarithm function is applied.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
log(threadgroup ST &dst, threadgroup const ST &src, const int threadIdx) {
|
||||
unary_map<base_ops::log, ST>(dst, src, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the absolute function to each element of the source tile and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the absolute function is applied.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
abs(threadgroup ST &dst, threadgroup const ST &src, const int threadIdx) {
|
||||
unary_map<base_ops::abs, ST>(dst, src, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the rectified linear unit function to each element of the source tile and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the rectified linear unit function is applied.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
relu(threadgroup ST &dst, threadgroup const ST &src, const int threadIdx) {
|
||||
unary_map<base_ops::relu, ST>(dst, src, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Copies the elements of the source tile to the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source data to be copied.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
copy(threadgroup ST &dst, thread const U &src, const int threadIdx) {
|
||||
bin_map<base_ops::copy, ST>(dst, src, threadIdx);
|
||||
}
|
||||
|
||||
// uniform binary maps
|
||||
/**
|
||||
* @brief Finds the maximum of each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile.
|
||||
* @param[in] rhs The second source data.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
max(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const int threadIdx) {
|
||||
bin_map<base_ops::max, ST>(dst, lhs, rhs, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Finds the minimum of each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile.
|
||||
* @param[in] rhs The second source data.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
min(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const int threadIdx) {
|
||||
bin_map<base_ops::min, ST>(dst, lhs, rhs, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Adds each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile.
|
||||
* @param[in] rhs The second source data.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
add(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const int threadIdx) {
|
||||
bin_map<base_ops::sum, ST>(dst, lhs, rhs, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Subtracts each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile.
|
||||
* @param[in] rhs The second source data.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
sub(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const int threadIdx) {
|
||||
bin_map<base_ops::sub, ST>(dst, lhs, rhs, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Multiplies each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile.
|
||||
* @param[in] rhs The second source data.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
mul(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const int threadIdx) {
|
||||
bin_map<base_ops::mul, ST>(dst, lhs, rhs, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Divides each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile.
|
||||
* @param[in] rhs The second source data.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
div(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const int threadIdx) {
|
||||
bin_map<base_ops::div, ST>(dst, lhs, rhs, threadIdx);
|
||||
}
|
||||
|
||||
// Row and col maps
|
||||
|
||||
/**
|
||||
* @brief Adds row values to each row of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the addition on.
|
||||
* @param row_values[in] Column vector containing values to add to each row.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
add_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const int threadIdx) {
|
||||
row_map<base_ops::sum, ST, SV>(dst, src, row_values, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Subtracts row values from each row of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the subtraction on.
|
||||
* @param row_values[in] Column vector containing values to subtract from each row.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
sub_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const int threadIdx) {
|
||||
row_map<base_ops::sub, ST, SV>(dst, src, row_values, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Multiplies each row of a tile by row values.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the multiplication on.
|
||||
* @param row_values[in] Column vector containing values to multiply each row by.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
mul_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const int threadIdx) {
|
||||
row_map<base_ops::mul, ST, SV>(dst, src, row_values, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Divides each row of a tile by row values.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the division on.
|
||||
* @param row_values[in] Column vector containing values to divide each row by.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
div_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const int threadIdx) {
|
||||
row_map<base_ops::div, ST, SV>(dst, src, row_values, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Broadcast a vector into into a tile's rows.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param row_values[in] Column vector containing values to broadcast into rows.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
broadcast_row(threadgroup ST &dst, threadgroup const SV &row_values, const int threadIdx) {
|
||||
row_map<base_ops::copy2, ST, SV>(dst, dst, row_values, threadIdx);
|
||||
}
|
||||
|
||||
|
||||
// col maps
|
||||
/**
|
||||
* @brief Adds column values to each column of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the addition on.
|
||||
* @param col_values[in] Row vector containing values to add to each column.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
add_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const int threadIdx) {
|
||||
col_map<base_ops::sum, ST, SV>(dst, src, col_values, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Subtracts column values from each column of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the subtraction on.
|
||||
* @param col_values[in] Row vector containing values to subtract from each column.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
sub_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const int threadIdx) {
|
||||
col_map<base_ops::sub, ST, SV>(dst, src, col_values, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Multiplies each column of a tile by column values.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the multiplication on.
|
||||
* @param col_values[in] Row vector containing values to multiply each column by.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
mul_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const int threadIdx) {
|
||||
col_map<base_ops::mul, ST, SV>(dst, src, col_values, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Divides each column of a tile by column values.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the division on.
|
||||
* @param col_values[in] Row vector containing values to divide each column by.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
div_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const int threadIdx) {
|
||||
col_map<base_ops::div, ST, SV>(dst, src, col_values, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Broadcast a vector into into a tile's columns.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param row_values[in] Row vector containing values to broadcast into cols.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
broadcast_col(threadgroup ST &dst, threadgroup const SV &col_values, const int threadIdx) {
|
||||
col_map<base_ops::copy2, ST, SV>(dst, dst, col_values, threadIdx);
|
||||
}
|
||||
@@ -0,0 +1,284 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Group reductions on shared tiles.
|
||||
*/
|
||||
|
||||
/**
|
||||
* Performs row-wise reduction on a matrix using a specified operation.
|
||||
*
|
||||
* @tparam op The operation to be applied for reduction.
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type with row layout.
|
||||
* @param row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param src The source matrix on which to perform the reduction.
|
||||
* @param src_accum The initial value of the accumulator, used when reset is false.
|
||||
* @param reset A boolean flag indicating whether to reset the accumulator (ignore src_accum) or not.
|
||||
*/
|
||||
template<typename op, typename SV, typename ST, bool reset>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_reduce(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
|
||||
using dtype = typename SV::dtype;
|
||||
for (int row = laneid(threadIdx); row < src.rows; row += GROUP_THREADS) {
|
||||
dtype accum = src[{row, 0}];
|
||||
#pragma clang loop unroll(full)
|
||||
for (int col = 1; col < src.cols; col++) {
|
||||
accum = op::template op<dtype>(accum, src[{row, col}]);
|
||||
}
|
||||
if (reset) {
|
||||
row_accum[row] = accum;
|
||||
} else {
|
||||
row_accum[row] = op::template op<dtype>(src_accum[row], accum);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs column-wise reduction on a matrix using a specified operation.
|
||||
*
|
||||
* @tparam op The operation to be applied for reduction.
|
||||
* @tparam V The shared vector type for the column accumulator.
|
||||
* @tparam T The shared matrix type with column layout.
|
||||
* @param col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param src The source matrix on which to perform the reduction.
|
||||
* @param src_accum The initial value of the accumulator, used when reset is false.
|
||||
* @param reset A boolean flag indicating whether to reset the accumulator (ignore src_accum) or not.
|
||||
*/
|
||||
template<typename op, typename SV, typename ST, bool reset>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_reduce(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
|
||||
using dtype = typename SV::dtype;
|
||||
for (int col = laneid(threadIdx); col < src.cols; col += GROUP_THREADS) {
|
||||
dtype accum = src[{0, col}];
|
||||
#pragma clang loop unroll(full)
|
||||
for (int row = 1; row < src.rows; row++) {
|
||||
accum = op::template op<dtype>(accum, src[{row, col}]);
|
||||
}
|
||||
if (reset) {
|
||||
col_accum[col] = accum;
|
||||
} else {
|
||||
col_accum[col] = op::template op<dtype>(src_accum[col], accum);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
|
||||
|
||||
/**
|
||||
* @brief Store the maximum of each row of the src shared matrix in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_max(threadgroup SV &row_accum, threadgroup const ST &src, const int threadIdx) {
|
||||
row_reduce<base_ops::max, SV, ST, true>(row_accum, src, row_accum, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Store the minimum of each row of the src shared matrix in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_min(threadgroup SV &row_accum, threadgroup const ST &src, const int threadIdx) {
|
||||
row_reduce<base_ops::min, SV, ST, true>(row_accum, src, row_accum, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Store the sum of each row of the src shared matrix in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_sum(threadgroup SV &row_accum, threadgroup const ST &src, const int threadIdx) {
|
||||
row_reduce<base_ops::sum, SV, ST, true>(row_accum, src, row_accum, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Store the product of each row of the src shared matrix in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_prod(threadgroup SV &row_accum, threadgroup const ST &src, const int threadIdx) {
|
||||
row_reduce<base_ops::mul, SV, ST, true>(row_accum, src, row_accum, threadIdx);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store the maximum of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_max(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
|
||||
row_reduce<base_ops::max, SV, ST, false>(row_accum, src, src_accum, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Store the minimum of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_min(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
|
||||
row_reduce<base_ops::min, SV, ST, false>(row_accum, src, src_accum, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Store the sum of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_sum(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
|
||||
row_reduce<base_ops::sum, SV, ST, false>(row_accum, src, src_accum, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Store the product of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_prod(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
|
||||
row_reduce<base_ops::mul, SV, ST, false>(row_accum, src, src_accum, threadIdx);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store the maximum of each column of the src shared matrix in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_max(threadgroup SV &col_accum, threadgroup const ST &src, const int threadIdx) {
|
||||
col_reduce<base_ops::max, SV, ST, true>(col_accum, src, col_accum, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Store the minimum of each column of the src shared matrix in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_min(threadgroup SV &col_accum, threadgroup const ST &src, const int threadIdx) {
|
||||
col_reduce<base_ops::min, threadgroup SV, threadgroup ST, true>(col_accum, src, col_accum, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Store the sum of each column of the src shared matrix in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_sum(threadgroup SV &col_accum, threadgroup const ST &src, const int threadIdx) {
|
||||
col_reduce<base_ops::sum, SV, ST, true>(col_accum, src, col_accum, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Store the product of each column of the src shared matrix in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_prod(threadgroup SV &col_accum, threadgroup const ST &src, const int threadIdx) {
|
||||
col_reduce<base_ops::mul, SV, ST, true>(col_accum, src, col_accum, threadIdx);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store the maximum of each column of the src shared matrix, as well as the src_accum shared vector, in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_max(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
|
||||
col_reduce<base_ops::max, SV, ST, false>(col_accum, src, src_accum, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Store the minimum of each column of the src shared matrix, as well as the src_accum shared vector, in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_min(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
|
||||
col_reduce<base_ops::min, SV, ST, false>(col_accum, src, src_accum, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Store the sum of each column of the src shared tile, as well as the src_accum row vector, in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_sum(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
|
||||
col_reduce<base_ops::sum, SV, ST, false>(col_accum, src, src_accum, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Store the product of each column of the src shared tile, as well as the src_accum row vector, in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_prod(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
|
||||
col_reduce<base_ops::mul, SV, ST, false>(col_accum, src, src_accum, threadIdx);
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
#include "conversions.metal"
|
||||
#include "maps.metal"
|
||||
#include "reductions.metal"
|
||||
@@ -0,0 +1,29 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Group conversions on shared vectors.
|
||||
*/
|
||||
|
||||
/**
|
||||
* @brief Copies data from one shared vector to another, converting data types if necessary.
|
||||
*
|
||||
* This function copies data from the source shared vector `src` to the destination shared vector `dst`.
|
||||
* If the data types of `src` and `dst` are the same, it performs a direct memory copy. Otherwise, it
|
||||
* converts each element from the source data type to the destination data type using the appropriate
|
||||
* converter before copying.
|
||||
*
|
||||
* @tparam SV1 The type of the destination shared vector, must satisfy the ducks::sv::all concept.
|
||||
* @tparam SV2 The type of the source shared vector, must satisfy the ducks::sv::all concept.
|
||||
* @param[out] dst The destination shared vector.
|
||||
* @param[in] src The source shared vector.
|
||||
* @note The lengths of `src` and `dst` must be equal. This is enforced at compile time.
|
||||
*/
|
||||
template<typename SV1, typename SV2>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV1>() && ducks::is_shared_vector<SV2>(), void>::type
|
||||
copy(threadgroup SV1 &dst, threadgroup const SV2 &src, const int threadIdx) {
|
||||
static_assert(SV1::length == SV2::length, "Source and destination vectors must have the same length.");
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid(threadIdx); i < dst.length; i+=GROUP_THREADS) {
|
||||
dst[i] = base_types::convertor<typename SV1::dtype, typename SV2::dtype>::convert(src[i]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,267 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Group maps on shared vectors.
|
||||
*/
|
||||
|
||||
/**
|
||||
* @brief Applies a unary operation to each element of a shared memory vector.
|
||||
*
|
||||
* @tparam op Unary operation type.
|
||||
* @tparam T Shared memory vector type.
|
||||
* @param dst[out] Destination vector in which to store the result.
|
||||
* @param src[in] Source vector to apply the unary operation.
|
||||
*/
|
||||
template<typename op, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
unary_op(threadgroup SV &dst, threadgroup const SV &src, const int threadIdx) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(auto cur = laneid(threadIdx); cur < SV::length; cur+=GROUP_THREADS) {
|
||||
dst[cur] = op::template op<typename SV::dtype>(src[cur]);
|
||||
}
|
||||
}
|
||||
/**
|
||||
* @brief Perform a binary operation on two shared vectors.
|
||||
*
|
||||
* @tparam op The binary operation to perform.
|
||||
* @tparam T The type of the vectors.
|
||||
* @param dst[out] The destination vector where the result is stored.
|
||||
* @param lhs[in] The left-hand side vector for the operation.
|
||||
* @param rhs[in] The right-hand side vector for the operation.
|
||||
*/
|
||||
template<typename op, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
bin_op(threadgroup SV &dst, threadgroup const SV &lhs, threadgroup const SV &rhs, const int threadIdx) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(auto cur = laneid(threadIdx); cur < SV::length; cur+=GROUP_THREADS) {
|
||||
dst[cur] = op::template op<typename SV::dtype>(lhs[cur], rhs[cur]);
|
||||
}
|
||||
}
|
||||
/**
|
||||
* @brief Perform a binary operation on a shared vector and a scalar.
|
||||
*
|
||||
* @tparam op The binary operation to perform.
|
||||
* @tparam T The type of the vector.
|
||||
* @param dst[out] The destination vector where the result is stored.
|
||||
* @param src[in] The source vector for the operation.
|
||||
* @param param[in] The scalar parameter for the operation.
|
||||
*/
|
||||
template<typename op, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
bin_op(threadgroup SV &dst, threadgroup const SV &src, thread const typename SV::dtype ¶m, const int threadIdx) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(auto cur = laneid(threadIdx); cur < SV::length; cur+=GROUP_THREADS) {
|
||||
dst[cur] = op::template op<typename SV::dtype>(src[cur], param);
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
|
||||
|
||||
// ---- const ops ----
|
||||
|
||||
/**
|
||||
* @brief Sets all elements of a shared memory vector to zero.
|
||||
*
|
||||
* @tparam T Shared memory vector type.
|
||||
* @param dst[out] Destination vector to be set to zero.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
zero(threadgroup SV &dst, const int threadIdx) {
|
||||
unary_op<base_ops::zero, SV>(dst, dst, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of a shared memory vector to one.
|
||||
*
|
||||
* @tparam T Shared memory vector type.
|
||||
* @param dst[out] Destination vector to be set to one.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
one(threadgroup SV &dst, const int threadIdx) {
|
||||
unary_op<base_ops::one, SV>(dst, dst, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of a shared memory vector to positive infinity.
|
||||
*
|
||||
* @tparam T Shared memory vector type.
|
||||
* @param dst[out] Destination vector to be set to positive infinity.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
pos_infty(threadgroup SV &dst, const int threadIdx) {
|
||||
unary_op<base_ops::pos_infty, SV>(dst, dst, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of a shared memory vector to negative infinity.
|
||||
*
|
||||
* @tparam T Shared memory vector type.
|
||||
* @param dst[out] Destination vector to be set to negative infinity.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
neg_infty(threadgroup SV &dst, const int threadIdx) {
|
||||
unary_op<base_ops::neg_infty, SV>(dst, dst, threadIdx);
|
||||
}
|
||||
|
||||
// ---- unary ops ----
|
||||
|
||||
/**
|
||||
* @brief Copies the elements from one shared vector to another.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the source vector.
|
||||
* @param dst[out] Destination vector where the elements will be copied to.
|
||||
* @param src[in] Source vector to copy the elements from.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
copy(threadgroup SV &dst, thread const U &src, const int threadIdx) {
|
||||
bin_op<base_ops::copy2, SV>(dst, dst, src, threadIdx); // the second arg is ignored here.
|
||||
}
|
||||
/**
|
||||
* @brief Applies the exponential function element-wise to a shared vector.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @param dst[out] Destination vector where the exponential values will be stored.
|
||||
* @param src[in] Source vector to apply the exponential function to.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
exp(threadgroup SV &dst, threadgroup const SV &src, const int threadIdx) {
|
||||
unary_op<base_ops::exp, SV>(dst, src, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the exponential function element-wise to a shared vector, in base 2.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @param dst[out] Destination vector where the exponential values will be stored.
|
||||
* @param src[in] Source vector to apply the exponential function to.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
exp2(threadgroup SV &dst, threadgroup const SV &src, const int threadIdx) {
|
||||
unary_op<base_ops::exp2, SV>(dst, src, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the natural logarithm function element-wise to a shared vector.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @param dst[out] Destination vector where the exponential values will be stored.
|
||||
* @param src[in] Source vector to apply the logarithm function to.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
log(threadgroup SV &dst, threadgroup const SV &src, const int threadIdx) {
|
||||
unary_op<base_ops::log, SV>(dst, src, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the absolute value function element-wise to a shared vector.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @param dst[out] Destination vector where the absolute values will be stored.
|
||||
* @param src[in] Source vector to apply the absolute value function to.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
abs(threadgroup SV &dst, threadgroup const SV &src, const int threadIdx) {
|
||||
unary_op<base_ops::abs, SV>(dst, src, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the rectified linear unit (ReLU) function element-wise to a shared vector.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @param dst[out] Destination vector where the ReLU values will be stored.
|
||||
* @param src[in] Source vector to apply the ReLU function to.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
relu(threadgroup SV &dst, threadgroup const SV &src, const int threadIdx) {
|
||||
unary_op<base_ops::relu, SV>(dst, src, threadIdx);
|
||||
}
|
||||
|
||||
// ---- binary ops ----
|
||||
|
||||
/**
|
||||
* @brief Computes the element-wise maximum of two shared vectors.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the maximum values will be stored.
|
||||
* @param lhs[in] First vector for the maximum operation.
|
||||
* @param rhs[in] Second vector for the maximum operation.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
max(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const int threadIdx) {
|
||||
bin_op<base_ops::max, SV>(dst, lhs, rhs, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise minimum of two shared vectors.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the minimum values will be stored.
|
||||
* @param lhs[in] First vector for the minimum operation.
|
||||
* @param rhs[in] Second vector for the minimum operation.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
min(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const int threadIdx) {
|
||||
bin_op<base_ops::min, SV>(dst, lhs, rhs, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise sum of two shared vectors.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the sum values will be stored.
|
||||
* @param lhs[in] First vector for the sum operation.
|
||||
* @param rhs[in] Second vector for the sum operation.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
add(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const int threadIdx) {
|
||||
bin_op<base_ops::sum, SV>(dst, lhs, rhs, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise difference of two shared vectors.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the difference values will be stored.
|
||||
* @param lhs[in] First vector for the difference operation.
|
||||
* @param rhs[in] Second vector for the difference operation.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
sub(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const int threadIdx) {
|
||||
bin_op<base_ops::sub, SV>(dst, lhs, rhs, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise product of two shared vectors.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the product values will be stored.
|
||||
* @param lhs[in] First vector for the product operation.
|
||||
* @param rhs[in] Second vector for the product operation.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
mul(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const int threadIdx) {
|
||||
bin_op<base_ops::mul, SV>(dst, lhs, rhs, threadIdx);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise division of two shared vectors.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the division values will be stored.
|
||||
* @param lhs[in] First vector for the division operation.
|
||||
* @param rhs[in] Second vector for the division operation.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
div(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const int threadIdx) {
|
||||
bin_op<base_ops::div, SV>(dst, lhs, rhs, threadIdx);
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
#include "conversions.metal"
|
||||
#include "maps.metal"
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
#pragma once
|
||||
#include "group/group.metal"
|
||||
#include "warp/warp.metal"
|
||||
@@ -0,0 +1,4 @@
|
||||
#pragma once
|
||||
#include "tile/tile.metal"
|
||||
#include "util/util.metal"
|
||||
#include "vec/vec.metal"
|
||||
@@ -0,0 +1,51 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Functions for transferring data directly between global memory and registers and back.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../../../../../common/common.metal"
|
||||
#include "../../../../../types/types.metal"
|
||||
|
||||
#include "../global_to_register.metal"
|
||||
|
||||
namespace mittens {
|
||||
/**
|
||||
* @brief Load data from source arrays into a complex-type tile.
|
||||
*
|
||||
* @tparam CRT The complex tile type.
|
||||
* @tparam U The data type of the source arrays.
|
||||
* @param dst[out] The destination tile to load data into.
|
||||
* @param resrc[in] The source array to load the real component data from.
|
||||
* @param imsrc[in] The source array to load the imaginary component data from.
|
||||
* @param re_row_stride[in] The stride in elements between rows in the real component source array.
|
||||
* @param im_row_stride[in] The stride in elements between rows in the imaginary component source array.
|
||||
*/
|
||||
template<typename CRT, typename CGL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_complex_register_tile<CRT>() && ducks::is_complex_global_layout<CGL>(), void>::type
|
||||
load(thread CRT &dst, thread const CGL &src, thread const coord &idx, const short laneid) {
|
||||
// Internally will use the correct load() method for row and column types
|
||||
load(dst.real, src.real, idx);
|
||||
load(dst.imag, src.imag, idx);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store data from a complex register tile to destination arrays in global memory.
|
||||
*
|
||||
* @tparam CRT The complex tile type.
|
||||
* @tparam U The data type of the destination arrays.
|
||||
* @param redst[out] The destination array in global memory to store the real component data into.
|
||||
* @param imdst[out] The destination array in global memory to store the imaginary component data into.
|
||||
* @param src[in] The source register tile to store data from.
|
||||
* @param re_row_stride[in] The stride in elements between rows in the real component destination array.
|
||||
* @param im_row_stride[in] The stride in elements between rows in the imaginary component destination array.
|
||||
*/
|
||||
template<typename CRT, typename CGL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_complex_register_tile<CRT>() && ducks::is_complex_global_layout<CGL>(), void>::type
|
||||
store(thread CGL &dst, thread const CRT &src, thread const coord &idx) {
|
||||
// Internally will use the correct load() method for row and column types
|
||||
store(dst.real, src.real, idx);
|
||||
store(dst.imag, src.imag, idx);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Functions for transferring data directly between global and shared memory and back.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../../../../../common/common.metal"
|
||||
#include "../../../../../types/types.metal"
|
||||
|
||||
#include "../global_to_shared.metal"
|
||||
|
||||
namespace mittens {
|
||||
/**
|
||||
* @brief Loads data from global memory into a complex shared memory tile with a row layout.
|
||||
*
|
||||
* @tparam CST The type of the complex shared tile.
|
||||
* @param[out] dst The destination complex shared memory tile.
|
||||
* @param[in] resrc The source global memory array for the real component.
|
||||
* @param[in] imsrc The source global memory array for the imaginary component.
|
||||
* @param re_row_stride[in] The stride between rows in the source real component array.
|
||||
* @param im_row_stride[in] The stride between rows in the source imaginary component array.
|
||||
*/
|
||||
template<typename CST, typename CGL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_complex_shared_tile<CST>() && ducks::is_global_layout<CGL>(), void>::type
|
||||
load(threadgroup CST &dst, thread const CGL &src, thread const coord &idx) {
|
||||
load(dst.real, src.real, idx);
|
||||
load(dst.imag, src.imag, idx);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Stores bf16 data from a complex shared memory tile with a row layout into global memory.
|
||||
*
|
||||
* @tparam CST The type of the complex shared tile.
|
||||
* @param[out] redst The destination global memory array for the real component.
|
||||
* @param[out] imdst The destination global memory array for the imaginary component.
|
||||
* @param[in] src The source complex shared memory tile.
|
||||
* @param re_row_stride[in] The stride between rows in the destination real component array.
|
||||
* @param im_row_stride[in] The stride between rows in the destination imaginary component array.
|
||||
*/
|
||||
template<typename CST, typename CGL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_complex_shared_tile<CST>() && ducks::is_complex_global_layout<CGL>(), void>::type
|
||||
store(thread const CGL &dst, threadgroup CST &src, thread const coord &idx) {
|
||||
store(dst.real, src.real, idx);
|
||||
store(dst.imag, src.imag, idx);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Functions for transferring data directly between shared memory and registers and back.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
|
||||
#include "../../../../../common/common.metal"
|
||||
#include "../../../../../types/types.metal"
|
||||
|
||||
#include "../shared_to_register.metal"
|
||||
|
||||
namespace mittens {
|
||||
/**
|
||||
* @brief Load data from a complex shared tile into a complex register tile.
|
||||
*
|
||||
* @tparam CRT The complex register tile type
|
||||
* @tparam CST The complex shared tile type
|
||||
* @param dst[out] The destination complex register tile.
|
||||
* @param src[in] The source complex shared tile.
|
||||
*/
|
||||
template<typename CRT, typename CST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_complex_shared_tile<CST>() && ducks::is_complex_register_tile<CRT>(), void>::type
|
||||
load(thread CRT &dst, threadgroup const CST &src) {
|
||||
load(dst.real, src.real);
|
||||
load(dst.imag, src.imag);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store data into a complex shared tile from a complex register tile.
|
||||
*
|
||||
* @tparam RT The complex register tile type
|
||||
* @tparam ST The complex shared tile type
|
||||
* @param dst[out] The destination complex shared tile.
|
||||
* @param src[in] The source complex register tile.
|
||||
*/
|
||||
template<typename CRT, typename CST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_complex_shared_tile<CST>() && ducks::is_complex_register_tile<CRT>(), void>::type
|
||||
store(threadgroup CST &dst, thread const CRT &src) {
|
||||
store(dst.real, src.real);
|
||||
store(dst.imag, src.imag);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
|
||||
@@ -0,0 +1,217 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Functions for transferring data directly between global memory and registers and back.
|
||||
*/
|
||||
|
||||
#pragma once // done!
|
||||
#include "../../../../types/types.metal"
|
||||
#include "../../../../common/common.metal"
|
||||
#include <metal_stdlib>
|
||||
namespace mittens{
|
||||
|
||||
namespace meta {
|
||||
template<typename RT, typename U>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>(), void>::type
|
||||
load(int i, int j, thread RT *dst, const device U *src_ptr, const short simd_y, const short simd_x, const int row_stride) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename RT::T2;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
using layout = typename RT::layout;
|
||||
unsigned offset = (simd_y + i * rt_base<T, layout>::tile_size) * row_stride + (simd_x + j * rt_base<T, layout>::tile_size);
|
||||
T2 src2 = base_types::convertor<T2, U2>::convert(*((device U2*)(&src_ptr[offset])));
|
||||
dst->tiles[i][j].data.thread_elements()[0] = src2[0];
|
||||
dst->tiles[i][j].data.thread_elements()[1] = src2[1];
|
||||
}
|
||||
|
||||
template<typename RT, typename U>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>(), void>::type
|
||||
load(int i, int j, thread RT *dst, const device U *src_ptr, const short simd_y, const short simd_x, const int row_stride) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename RT::T2;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
using layout = typename RT::layout;
|
||||
unsigned offset = (simd_y + i * rt_base<T, layout>::tile_size) * row_stride + (simd_x + j * rt_base<T, layout>::tile_size);
|
||||
dst->tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src_ptr[offset]);
|
||||
offset += row_stride;
|
||||
dst->tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src_ptr[offset]);
|
||||
}
|
||||
|
||||
template<typename RT, typename U>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>(), void>::type
|
||||
store(int i, int j, device U *dst_ptr, const thread RT *src, const short simd_y, const short simd_x, const int row_stride) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename RT::T2;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
using layout = typename RT::layout;
|
||||
unsigned offset = (simd_y + i * TILE_DIM) * row_stride + (simd_x + j * TILE_DIM);
|
||||
U2 src2 = base_types::convertor<U2, T2>::convert(
|
||||
T2(src->tiles[i][j].data.thread_elements()[0],
|
||||
src->tiles[i][j].data.thread_elements()[1])
|
||||
);
|
||||
*((device U2*)&dst_ptr[offset]) = src2;
|
||||
}
|
||||
|
||||
template<typename RT, typename U>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>(), void>::type
|
||||
store(int i, int j, device U *dst_ptr, const thread RT *src, const short simd_y, const short simd_x, const int row_stride) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename RT::T2;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
using layout = typename RT::layout;
|
||||
unsigned offset = (simd_y + i * rt_base<T, layout>::tile_size) * row_stride + (simd_x + j * rt_base<T, layout>::tile_size);
|
||||
dst_ptr[offset] = base_types::convertor<U, T>::convert(src->tiles[i][j].data.thread_elements()[0]);
|
||||
offset += row_stride;
|
||||
dst_ptr[offset] = base_types::convertor<U, T>::convert(src->tiles[i][j].data.thread_elements()[1]);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Load data from a source array into a row-major layout tile.
|
||||
*
|
||||
* @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<typename RT, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
|
||||
load(thread RT &dst, thread const GL &src, thread const coord &idx, const short laneid) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename RT::T2;
|
||||
using U = typename GL::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
using layout = typename RT::layout;
|
||||
const device U *src_ptr = (device U*)&src.template get<RT>(idx);
|
||||
const int row_stride = src.row_stride();
|
||||
|
||||
const short qid = laneid / 4;
|
||||
const short simd_y = (qid & 4) + (laneid / 2) % 4;
|
||||
const short simd_x = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
|
||||
// #pragma clang loop unroll(full)
|
||||
// for (int i = 0; i < RT::height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for (int j = 0; j < RT::width; j++) {
|
||||
// unsigned offset = (simd_y + i * rt_base<T, layout>::tile_size) * row_stride + (simd_x + j * rt_base<T, layout>::tile_size);
|
||||
// T2 src2 = base_types::convertor<T2, U2>::convert(*((device U2*)(&src_ptr[offset])));
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = src2[0];
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = src2[1];
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::load<RT, U>, &dst, src_ptr, simd_y, simd_x, row_stride);
|
||||
}
|
||||
/**
|
||||
* @brief Load data from a source array into a col-major layout tile.
|
||||
*
|
||||
* @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<typename RT, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
|
||||
load(thread RT &dst, thread const GL &src, thread const coord &idx, const short laneid) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename RT::T2;
|
||||
using U = typename GL::dtype;
|
||||
using layout = typename RT::layout;
|
||||
const device U *src_ptr = (device U*)&(src.template get<RT>(idx));
|
||||
const int row_stride = src.row_stride();
|
||||
|
||||
const short qid = laneid / 4;
|
||||
const short simd_x = (qid & 4) + (laneid / 2) % 4;
|
||||
const short simd_y = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
|
||||
// #pragma clang loop unroll(full)
|
||||
// for (int i = 0; i < RT::height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for (int j = 0; j < RT::width; j++) {
|
||||
// unsigned offset = (simd_y + i * rt_base<T, layout>::tile_size) * row_stride + (simd_x + j * rt_base<T, layout>::tile_size);
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src_ptr[offset]);
|
||||
// offset += row_stride;
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src_ptr[offset]);
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::load<RT, U>, &dst, src_ptr, simd_y, simd_x, row_stride);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store data from a register tile 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<typename RT, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
|
||||
store(thread GL &dst, thread const RT &src, thread const coord &idx, const short laneid) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename RT::T2;
|
||||
using U = typename GL::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
using layout = typename RT::layout;
|
||||
device U *dst_ptr = (device U*)&(dst.template get<RT>(idx));
|
||||
// device U* dst_ptr = dst.raw_ptr;
|
||||
const int row_stride = dst.row_stride();
|
||||
const short qid = laneid / 4;
|
||||
const short simd_y = (qid & 4) + (laneid / 2) % 4;
|
||||
const short simd_x = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
|
||||
// #pragma clang loop unroll(full)
|
||||
// for (int i = 0; i < RT::height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for (int j = 0; j < RT::width; j++) {
|
||||
// unsigned offset = (simd_y + i * TILE_DIM) * row_stride + (simd_x + j * TILE_DIM);
|
||||
// U2 src2 = base_types::convertor<U2, T2>::convert(
|
||||
// T2(src.tiles[i][j].data.thread_elements()[0],
|
||||
// src.tiles[i][j].data.thread_elements()[1])
|
||||
// );
|
||||
// *((device U2*)&dst_ptr[offset]) = src2;
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::store<RT, U>, dst_ptr, &src, simd_y, simd_x, row_stride);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store data from a register tile to a destination array in global memory with a col-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<typename RT, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
|
||||
store(thread GL &dst, thread const RT &src, thread const coord &idx, const short laneid) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename RT::T2;
|
||||
using U = typename GL::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
using layout = typename RT::layout;
|
||||
device U *dst_ptr = (device U*)&(dst.template get<RT>(idx));
|
||||
const int row_stride = dst.row_stride();
|
||||
const short qid = laneid / 4;
|
||||
const short simd_x = (qid & 4) + (laneid / 2) % 4;
|
||||
const short simd_y = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
|
||||
// #pragma clang loop unroll(full)
|
||||
// for (int i = 0; i < RT::height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for (int j = 0; j < RT::width; j++) {
|
||||
// unsigned offset = (simd_y + i * rt_base<T, layout>::tile_size) * row_stride + (simd_x + j * rt_base<T, layout>::tile_size);
|
||||
// dst_ptr[offset] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[0]);
|
||||
// offset += row_stride;
|
||||
// dst_ptr[offset] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[1]);
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::store<RT, U>, dst_ptr, &src, simd_y, simd_x, row_stride);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Functions for transferring data directly between global and shared memory and back.
|
||||
*/
|
||||
|
||||
#pragma once // not done!
|
||||
#include "../../../../types/types.metal"
|
||||
#include "../../../../common/common.metal"
|
||||
#include <metal_stdlib>
|
||||
namespace mittens {
|
||||
|
||||
//
|
||||
namespace meta {
|
||||
template<typename ST, int memcpy_per_row, int elem_per_memcpy, int READ_FLOATS>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
load(int i, threadgroup ST *dst, device const typename ST::dtype *src, thread const int& row_stride, thread const short& laneid) {
|
||||
{
|
||||
unsigned idx = i + laneid;
|
||||
unsigned row = idx / memcpy_per_row;
|
||||
unsigned col = (idx*elem_per_memcpy) % ST::cols;
|
||||
*(threadgroup ReadVector<READ_FLOATS>*)(&(*dst)[int2(row, col)]) = *(device ReadVector<READ_FLOATS>*)(&src[row*row_stride + col]);
|
||||
}
|
||||
}
|
||||
|
||||
template<typename ST, int memcpy_per_row, int elem_per_memcpy, int READ_FLOATS>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
store(int i, device typename ST::dtype *dst, threadgroup const ST *src, thread const int& row_stride, thread const short& laneid) {
|
||||
{
|
||||
unsigned idx = i + laneid;
|
||||
unsigned row = idx / memcpy_per_row;
|
||||
unsigned col = (idx*elem_per_memcpy) % ST::cols;
|
||||
*(device ReadVector<READ_FLOATS>*)(&dst[row*row_stride + col]) = *(threadgroup ReadVector<READ_FLOATS>*)(&(*src)[int2(row, col)]);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace meta
|
||||
|
||||
//
|
||||
///**
|
||||
// * @brief Loads data from global memory into a shared memory tile with a row layout.
|
||||
// *
|
||||
// * @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 row_stride[in] The stride between rows in the source array.
|
||||
// * @param laneid[in] Thread's index in SIMD group
|
||||
// */
|
||||
//template<typename ST>
|
||||
//static METAL_FUNC void load(threadgroup ST &dst, device const typename ST::dtype *src, const int row_stride, short laneid) {
|
||||
// using read_type = float;
|
||||
// ducks::assert_shared_tile<ST>();
|
||||
// constexpr const unsigned elem_per_memcpy = sizeof(read_type)/sizeof(typename ST::dtype); // 2
|
||||
// constexpr const unsigned memcpy_per_row = ST::cols / elem_per_memcpy; // 32/2=16 not power of 2
|
||||
// constexpr const unsigned total_calls = ST::num_elements / (SIMD_THREADS*elem_per_memcpy); // 1024/(32*2)=16
|
||||
//// #pragma clang loop unroll_count(1)
|
||||
//// #pragma clang loop unroll(disable)
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(unsigned i = 0; i < total_calls; i++) {
|
||||
// unsigned idx = i * 32 + laneid;
|
||||
// unsigned row = idx / memcpy_per_row;
|
||||
// unsigned col = (idx*elem_per_memcpy) % ST::cols;
|
||||
// *(threadgroup read_type*)(&dst[int2(row, col)]) = *(device read_type*)(&src[row*row_stride + col]);
|
||||
// }
|
||||
//
|
||||
//// ducks::assert_shared_tile<ST>();
|
||||
//// const constexpr int read_size = 1;
|
||||
//// using read_type = ReadVector<read_size>;
|
||||
//// constexpr const unsigned elem_per_memcpy = sizeof(read_type)/sizeof(typename ST::dtype); // 2
|
||||
//// constexpr const unsigned memcpy_per_row = ST::cols / elem_per_memcpy; // 32/2=16 not power of 2
|
||||
//// constexpr const unsigned total_calls = ST::num_elements / (SIMD_THREADS*elem_per_memcpy); // 1024/(32*2)=16
|
||||
////
|
||||
////
|
||||
//// meta::unroll_i_in_range<0, total_calls * SIMD_THREADS, SIMD_THREADS>::run(meta::load<ST, memcpy_per_row, elem_per_memcpy, read_size>, &dst, src, row_stride, laneid);
|
||||
//}
|
||||
//
|
||||
//
|
||||
///**
|
||||
// * @brief Stores data from a shared memory tile with a row layout 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.
|
||||
// * @param laneid[in] Thread's index in SIMD group
|
||||
// */
|
||||
//template<typename ST>
|
||||
//static METAL_FUNC void store(device typename ST::dtype *dst, threadgroup const ST &src, const int row_stride, short laneid) {
|
||||
// using read_type = float4;
|
||||
// ducks::assert_shared_tile<ST>();
|
||||
// constexpr const unsigned elem_per_memcpy = sizeof(read_type)/sizeof(typename ST::dtype);
|
||||
// constexpr const unsigned memcpy_per_row = ST::cols / elem_per_memcpy;
|
||||
// constexpr const unsigned total_calls = ST::num_elements / (SIMD_THREADS*elem_per_memcpy);
|
||||
//// #pragma clang loop unroll_count(READ_SIZE)
|
||||
////#pragma clang loop unroll(disable)
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(unsigned i = 0; i < total_calls; i++) {
|
||||
// unsigned idx = i * 32 + laneid;
|
||||
// unsigned row = idx / memcpy_per_row;
|
||||
// unsigned col = (idx*elem_per_memcpy) % src.cols;
|
||||
// *(device read_type*)(&dst[row*row_stride + col]) = *(threadgroup read_type*)(&src[int2(row, col)]);
|
||||
// }
|
||||
//
|
||||
////
|
||||
//// ducks::assert_shared_tile<ST>();
|
||||
//// const constexpr int read_size = 1;
|
||||
//// using read_type = ReadVector<read_size>;
|
||||
////
|
||||
//// constexpr const unsigned elem_per_memcpy = sizeof(read_type)/sizeof(typename ST::dtype);
|
||||
//// constexpr const unsigned memcpy_per_row = ST::cols / elem_per_memcpy;
|
||||
//// constexpr const unsigned total_calls = ST::num_elements / (SIMD_THREADS*elem_per_memcpy);
|
||||
////
|
||||
////
|
||||
//// meta::unroll_i_in_range<0, total_calls * SIMD_THREADS, SIMD_THREADS>::run(meta::store<ST, memcpy_per_row, elem_per_memcpy, read_size>, dst, &src, row_stride, laneid);
|
||||
//}
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* @brief Loads data from global memory into a shared memory tile with a row layout.
|
||||
*
|
||||
* @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 row_stride[in] The stride between rows in the source array.
|
||||
* @param laneid[in] Thread's index in SIMD group
|
||||
*/
|
||||
template<typename ST, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
|
||||
load(threadgroup ST &dst, thread const GL &src, thread const coord &idx, short laneid) {
|
||||
using U = typename GL::dtype;
|
||||
constexpr const int read_size = 1;
|
||||
using read_type = ReadVector<read_size>;
|
||||
device U *src_ptr = (device U*)&src.template get<ST>(idx);
|
||||
const int row_stride = src.row_stride();
|
||||
constexpr const unsigned elem_per_memcpy = sizeof(read_type)/sizeof(typename ST::dtype); // 2
|
||||
constexpr const unsigned memcpy_per_row = ST::cols / elem_per_memcpy; // 32/2=16 not power of 2
|
||||
constexpr const unsigned total_calls = ST::num_elements / (SIMD_THREADS*elem_per_memcpy); // 1024/(32*2)=16
|
||||
// #pragma clang loop unroll_count(1)
|
||||
// #pragma clang loop unroll(disable)
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(unsigned i = 0; i < total_calls; i++) {
|
||||
// unsigned idx = i * 32 + laneid;
|
||||
// unsigned row = idx / memcpy_per_row;
|
||||
// unsigned col = (idx*elem_per_memcpy) % ST::cols;
|
||||
// *(threadgroup read_type*)(&dst[int2(row, col)]) = *(device read_type*)(&src_ptr[row*row_stride + col]);
|
||||
// }
|
||||
meta::unroll_i_in_range<0, total_calls * SIMD_THREADS, SIMD_THREADS>::run(meta::load<ST, memcpy_per_row, elem_per_memcpy, read_size>, &dst, src_ptr, row_stride, laneid);
|
||||
}
|
||||
/*
|
||||
|
||||
*/
|
||||
|
||||
|
||||
/**
|
||||
* @brief Stores data from a shared memory tile with a row layout 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.
|
||||
* @param laneid[in] Thread's index in SIMD group
|
||||
*/
|
||||
template<typename ST, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
|
||||
store(thread GL &dst, threadgroup const ST &src, thread const coord &idx, short laneid) {
|
||||
using U = typename GL::dtype;
|
||||
constexpr const int read_size = 1;
|
||||
using read_type = ReadVector<read_size>;
|
||||
device U *dst_ptr = (device U*)&dst.template get<ST>(idx);
|
||||
const int row_stride = dst.row_stride();
|
||||
|
||||
constexpr const unsigned elem_per_memcpy = sizeof(read_type)/sizeof(typename ST::dtype);
|
||||
constexpr const unsigned memcpy_per_row = ST::cols / elem_per_memcpy;
|
||||
constexpr const unsigned total_calls = ST::num_elements / (SIMD_THREADS*elem_per_memcpy);
|
||||
// #pragma clang loop unroll_count(READ_SIZE)
|
||||
//#pragma clang loop unroll(disable)
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(unsigned i = 0; i < total_calls; i++) {
|
||||
// unsigned idx = i * 32 + laneid;
|
||||
// unsigned row = idx / memcpy_per_row;
|
||||
// unsigned col = (idx*elem_per_memcpy) % src.cols;
|
||||
// *(device read_type*)(&dst_ptr[row*row_stride + col]) = *(threadgroup read_type*)(&src[int2(row, col)]);
|
||||
// }
|
||||
|
||||
meta::unroll_i_in_range<0, total_calls * SIMD_THREADS, SIMD_THREADS>::run(meta::store<ST, memcpy_per_row, elem_per_memcpy, read_size>, dst_ptr, &src, row_stride, laneid);
|
||||
}
|
||||
|
||||
|
||||
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,461 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Functions for transferring data directly between shared memory and registers and back.
|
||||
*/
|
||||
#pragma once // done!
|
||||
|
||||
#include "../../../../types/types.metal"
|
||||
#include "../../../../common/common.metal"
|
||||
#include <metal_stdlib>
|
||||
namespace mittens {
|
||||
|
||||
// These probably need to be redone to reduce bank conflicts.
|
||||
// They currently work fine with xor layout but it should be
|
||||
// possible to reduce their bank conflicts with other layouts too.
|
||||
//
|
||||
namespace meta {
|
||||
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
loadStR(int i, int j, thread RT *dst, threadgroup const ST *src, short laneid, int offsetY, int offsetX) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
int y = offsetY + i * mittens::TILE_DIM;
|
||||
int x = offsetX + j * mittens::TILE_DIM;
|
||||
T2 values = base_types::convertor<T2, U2>::convert(*((threadgroup U2*)(&(*src)[int2(y, x)])));
|
||||
dst->tiles[i][j].data.thread_elements()[0] = values[0];
|
||||
dst->tiles[i][j].data.thread_elements()[1] = values[1];
|
||||
//
|
||||
// simdgroup_load(dst->tiles[i][j].data,
|
||||
// (threadgroup T*)(src->data),
|
||||
// src->cols,
|
||||
// {i * mittens::TILE_DIM, j * mittens::TILE_DIM},
|
||||
//
|
||||
}
|
||||
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
storeStR(int i, int j, threadgroup ST *dst, thread const RT *src, short laneid, int offsetY, int offsetX) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
int y = offsetY + i * mittens::TILE_DIM;
|
||||
int x = offsetX + j * mittens::TILE_DIM;
|
||||
U2 values = base_types::convertor<U2, T2>::convert({src->tiles[i][j].data.thread_elements()[0], src->tiles[i][j].data.thread_elements()[1]});
|
||||
*((threadgroup U2*)(&(*dst)[int2(y, x)])) = values;
|
||||
}
|
||||
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
loadStR(int i, int j, thread RT *dst, threadgroup const ST *src, short laneid, int offsetY, int offsetX) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
int y = offsetY + i * mittens::TILE_DIM;
|
||||
int x = offsetX + j * mittens::TILE_DIM;
|
||||
// dst->tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert((*src)[int2(y , x)]);
|
||||
// dst->tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert((*src)[int2(y+1, x)]);
|
||||
T2 vals = base_types::convertor<T2, U2>::convert({(*src)[int2(y , x)], (*src)[int2(y+1, x)]});
|
||||
dst->tiles[i][j].data.thread_elements()[0] = vals[0];
|
||||
dst->tiles[i][j].data.thread_elements()[1] = vals[1];
|
||||
}
|
||||
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
storeStR(int i, int j, threadgroup ST *dst, thread const RT *src, short laneid, int offsetY, int offsetX) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
int y = offsetY + i * mittens::TILE_DIM;
|
||||
int x = offsetX + j * mittens::TILE_DIM;
|
||||
// (*dst)[int2(y , x)] = base_types::convertor<U, T>::convert(src->tiles[i][j].data.thread_elements()[0]);
|
||||
// (*dst)[int2(y+1, x)] = base_types::convertor<U, T>::convert(src->tiles[i][j].data.thread_elements()[1]);
|
||||
|
||||
U2 vals = base_types::convertor<U2, T2>::convert({src->tiles[i][j].data.thread_elements()[0], src->tiles[i][j].data.thread_elements()[1]});
|
||||
(*dst)[int2(y , x)] = vals[0];
|
||||
(*dst)[int2(y+1, x)] = vals[1];
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Load data from a shared tile into a register tile.
|
||||
*
|
||||
* @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.
|
||||
* @param laneid[in] Thread's index in SIMD group
|
||||
*/
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
load(thread RT &dst, threadgroup const ST &src, short laneid) {
|
||||
static_assert(RT::height == ST::height, "register tile and shared tile must match height");
|
||||
static_assert(RT::width == ST::width, "register tile and shared tile must match width");
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
const short qid = laneid / 4;
|
||||
int offsetY = (qid & 4) + (laneid / 2) % 4;
|
||||
int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < dst.height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < dst.width; j++) {
|
||||
// int y = offsetY + i * mittens::TILE_DIM;
|
||||
// int x = offsetX + j * mittens::TILE_DIM;
|
||||
// T2 values = base_types::convertor<T2, U2>::convert(*((threadgroup U2*)(&src[int2(y, x)])));
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = values[0];
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = values[1];
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::loadStR<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Load data from a shared tile into a register tile.
|
||||
*
|
||||
* @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.
|
||||
* @param laneid[in] Thread's index in SIMD group
|
||||
*/
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
load(thread RT &dst, threadgroup const ST &src, short laneid) {
|
||||
static_assert(RT::height == ST::height, "register tile and shared tile must match height");
|
||||
static_assert(RT::width == ST::width, "register tile and shared tile must match width");
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
const short qid = laneid / 4;
|
||||
// int offsetY = (qid & 4) + (laneid / 2) % 4;
|
||||
// int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
int offsetX = (qid & 4) + (laneid / 2) % 4;
|
||||
int offsetY = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < dst.height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < dst.width; j++) {
|
||||
// int y = offsetY + i * mittens::TILE_DIM;
|
||||
// int x = offsetX + j * mittens::TILE_DIM;
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src[int2(y , x)]);
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src[int2(y+1, x)]);
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::loadStR<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store data into a shared tile from a register tile.
|
||||
*
|
||||
* @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.
|
||||
* @param laneid[in] Thread's index in SIMD group
|
||||
*/
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
store(threadgroup ST &dst, thread const RT &src, short laneid) {
|
||||
ducks::assert_register_tile<RT>();
|
||||
ducks::assert_shared_tile<ST>();
|
||||
static_assert(RT::height == ST::height, "register tile and shared tile must match height");
|
||||
static_assert(RT::width == ST::width, "register tile and shared tile must match width");
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
|
||||
const short qid = laneid / 4;
|
||||
int offsetY = (qid & 4) + (laneid / 2) % 4;
|
||||
int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < src.height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < src.width; j++) {
|
||||
// int y = offsetY + i * mittens::TILE_DIM;
|
||||
// int x = offsetX + j * mittens::TILE_DIM;
|
||||
// U2 values = base_types::convertor<U2, T2>::convert({src.tiles[i][j].data.thread_elements()[0], src.tiles[i][j].data.thread_elements()[1]});
|
||||
// *((threadgroup U2*)(&dst[int2(y, x)])) = values;
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::storeStR<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store data into a shared tile from a register tile.
|
||||
*
|
||||
* @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.
|
||||
* @param laneid[in] Thread's index in SIMD group
|
||||
*/
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
store(threadgroup ST &dst, thread const RT &src, short laneid) {
|
||||
ducks::assert_register_tile<RT>();
|
||||
ducks::assert_shared_tile<ST>();
|
||||
static_assert(RT::height == ST::height, "register tile and shared tile must match height");
|
||||
static_assert(RT::width == ST::width, "register tile and shared tile must match width");
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
|
||||
const short qid = laneid / 4;
|
||||
// int offsetY = (qid & 4) + (laneid / 2) % 4;
|
||||
// int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
int offsetX = (qid & 4) + (laneid / 2) % 4;
|
||||
int offsetY = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < src.height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < src.width; j++) {
|
||||
// int y = offsetY + i * mittens::TILE_DIM;
|
||||
// int x = offsetX + j * mittens::TILE_DIM;
|
||||
// dst[int2(y , x)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[0]);
|
||||
// dst[int2(y+1, x)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[1]);
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::storeStR<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
|
||||
}
|
||||
|
||||
/*---------------------------------------------------------------------------------*/
|
||||
// These probably need to be redone to reduce bank conflicts.
|
||||
// They currently work fine with xor layout but it should be
|
||||
// possible to reduce their bank conflicts with other layouts too.
|
||||
//
|
||||
namespace meta {
|
||||
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
loadStR_r(int i, int j, thread RT *dst, thread const ST *src, short laneid, int offsetY, int offsetX) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
int y = offsetY + i * mittens::TILE_DIM;
|
||||
int x = offsetX + j * mittens::TILE_DIM;
|
||||
T2 values = base_types::convertor<T2, U2>::convert(*((threadgroup U2*)(&(*src)[int2(y, x)])));
|
||||
dst->tiles[i][j].data.thread_elements()[0] = values[0];
|
||||
dst->tiles[i][j].data.thread_elements()[1] = values[1];
|
||||
//
|
||||
// simdgroup_load(dst->tiles[i][j].data,
|
||||
// (threadgroup T*)(src->data),
|
||||
// src->cols,
|
||||
// {i * mittens::TILE_DIM, j * mittens::TILE_DIM},
|
||||
//
|
||||
}
|
||||
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
storeStR_r(int i, int j, thread ST *dst, thread const RT *src, short laneid, int offsetY, int offsetX) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
int y = offsetY + i * mittens::TILE_DIM;
|
||||
int x = offsetX + j * mittens::TILE_DIM;
|
||||
U2 values = base_types::convertor<U2, T2>::convert({src->tiles[i][j].data.thread_elements()[0], src->tiles[i][j].data.thread_elements()[1]});
|
||||
*((threadgroup U2*)(&(*dst)[int2(y, x)])) = values;
|
||||
}
|
||||
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
loadStR_c(int i, int j, thread RT *dst, thread const ST *src, short laneid, int offsetY, int offsetX) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
int y = offsetY + i * mittens::TILE_DIM;
|
||||
int x = offsetX + j * mittens::TILE_DIM;
|
||||
// dst->tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert((*src)[int2(y , x)]);
|
||||
// dst->tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert((*src)[int2(y+1, x)]);
|
||||
T2 vals = base_types::convertor<T2, U2>::convert({(*src)[int2(y , x)], (*src)[int2(y+1, x)]});
|
||||
dst->tiles[i][j].data.thread_elements()[0] = vals[0];
|
||||
dst->tiles[i][j].data.thread_elements()[1] = vals[1];
|
||||
}
|
||||
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
storeStR_c(int i, int j, thread ST *dst, thread const RT *src, short laneid, int offsetY, int offsetX) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
int y = offsetY + i * mittens::TILE_DIM;
|
||||
int x = offsetX + j * mittens::TILE_DIM;
|
||||
// (*dst)[int2(y , x)] = base_types::convertor<U, T>::convert(src->tiles[i][j].data.thread_elements()[0]);
|
||||
// (*dst)[int2(y+1, x)] = base_types::convertor<U, T>::convert(src->tiles[i][j].data.thread_elements()[1]);
|
||||
|
||||
U2 vals = base_types::convertor<U2, T2>::convert({src->tiles[i][j].data.thread_elements()[0], src->tiles[i][j].data.thread_elements()[1]});
|
||||
(*dst)[int2(y , x)] = vals[0];
|
||||
(*dst)[int2(y+1, x)] = vals[1];
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Load data from a shared tile into a register tile.
|
||||
*
|
||||
* @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.
|
||||
* @param laneid[in] Thread's index in SIMD group
|
||||
*/
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
load(thread RT &dst, thread const ST &src, short laneid) {
|
||||
static_assert(RT::height == ST::height, "register tile and shared tile must match height");
|
||||
static_assert(RT::width == ST::width, "register tile and shared tile must match width");
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
const short qid = laneid / 4;
|
||||
int offsetY = (qid & 4) + (laneid / 2) % 4;
|
||||
int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < dst.height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < dst.width; j++) {
|
||||
// int y = offsetY + i * mittens::TILE_DIM;
|
||||
// int x = offsetX + j * mittens::TILE_DIM;
|
||||
// T2 values = base_types::convertor<T2, U2>::convert(*((threadgroup U2*)(&src[int2(y, x)])));
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = values[0];
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = values[1];
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::loadStR_r<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Load data from a shared tile into a register tile.
|
||||
*
|
||||
* @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.
|
||||
* @param laneid[in] Thread's index in SIMD group
|
||||
*/
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
load(thread RT &dst, thread const ST &src, short laneid) {
|
||||
static_assert(RT::height == ST::height, "register tile and shared tile must match height");
|
||||
static_assert(RT::width == ST::width, "register tile and shared tile must match width");
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
const short qid = laneid / 4;
|
||||
// int offsetY = (qid & 4) + (laneid / 2) % 4;
|
||||
// int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
int offsetX = (qid & 4) + (laneid / 2) % 4;
|
||||
int offsetY = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < dst.height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < dst.width; j++) {
|
||||
// int y = offsetY + i * mittens::TILE_DIM;
|
||||
// int x = offsetX + j * mittens::TILE_DIM;
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src[int2(y , x)]);
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src[int2(y+1, x)]);
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::loadStR_c<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store data into a shared tile from a register tile.
|
||||
*
|
||||
* @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.
|
||||
* @param laneid[in] Thread's index in SIMD group
|
||||
*/
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
store(thread ST &dst, thread const RT &src, short laneid) {
|
||||
ducks::assert_register_tile<RT>();
|
||||
ducks::assert_shared_tile<ST>();
|
||||
static_assert(RT::height == ST::height, "register tile and shared tile must match height");
|
||||
static_assert(RT::width == ST::width, "register tile and shared tile must match width");
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
|
||||
const short qid = laneid / 4;
|
||||
int offsetY = (qid & 4) + (laneid / 2) % 4;
|
||||
int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < src.height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < src.width; j++) {
|
||||
// int y = offsetY + i * mittens::TILE_DIM;
|
||||
// int x = offsetX + j * mittens::TILE_DIM;
|
||||
// U2 values = base_types::convertor<U2, T2>::convert({src.tiles[i][j].data.thread_elements()[0], src.tiles[i][j].data.thread_elements()[1]});
|
||||
// *((threadgroup U2*)(&dst[int2(y, x)])) = values;
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::storeStR_r<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store data into a shared tile from a register tile.
|
||||
*
|
||||
* @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.
|
||||
* @param laneid[in] Thread's index in SIMD group
|
||||
*/
|
||||
template<typename RT, typename ST>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
|
||||
store(thread ST &dst, thread const RT &src, short laneid) {
|
||||
ducks::assert_register_tile<RT>();
|
||||
ducks::assert_shared_tile<ST>();
|
||||
static_assert(RT::height == ST::height, "register tile and shared tile must match height");
|
||||
static_assert(RT::width == ST::width, "register tile and shared tile must match width");
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
using U = typename ST::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
|
||||
const short qid = laneid / 4;
|
||||
// int offsetY = (qid & 4) + (laneid / 2) % 4;
|
||||
// int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
int offsetX = (qid & 4) + (laneid / 2) % 4;
|
||||
int offsetY = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < src.height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < src.width; j++) {
|
||||
// int y = offsetY + i * mittens::TILE_DIM;
|
||||
// int x = offsetX + j * mittens::TILE_DIM;
|
||||
// dst[int2(y , x)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[0]);
|
||||
// dst[int2(y+1, x)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[1]);
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::storeStR_c<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
#pragma once
|
||||
|
||||
#include "global_to_register.metal"
|
||||
#include "global_to_shared.metal"
|
||||
#include "shared_to_register.metal"
|
||||
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief General utilities not specialized for either tiles or vectors.
|
||||
*/
|
||||
#pragma once // done!
|
||||
#include "../tile/tile.metal"
|
||||
#include "../../../../types/shared/shared.metal"
|
||||
namespace mittens {
|
||||
|
||||
// sizeof() can be unreliable when working with references to objects
|
||||
// plus, template magic allows arrays of these objects to be copied, too.
|
||||
namespace detail {
|
||||
|
||||
template <typename T, uint32_t... dims>
|
||||
struct size_info;
|
||||
|
||||
template <typename T>
|
||||
struct size_info<T> {
|
||||
private:
|
||||
static_assert(ducks::is_shared_tile<T>() || ducks::is_shared_vector<T>(), "T must be a shared tile or shared vector");
|
||||
constant static constexpr uint32_t elements = ducks::is_shared_tile<T>() ? T::num_elements : T::length;
|
||||
constant static constexpr uint32_t bytes = elements * sizeof(typename T::dtype);
|
||||
};
|
||||
|
||||
template <typename T, uint32_t dim, uint32_t... rest_dims>
|
||||
struct size_info<T, dim, rest_dims...> {
|
||||
constant static constexpr uint32_t elements = dim * size_info<T, rest_dims...>::elements;
|
||||
constant static constexpr uint32_t bytes = dim * size_info<T, rest_dims...>::bytes;
|
||||
};
|
||||
}
|
||||
|
||||
template<typename T, uint32_t... dims> constant constexpr uint32_t size_elements = detail::size_info<T, dims...>::elements;
|
||||
template<typename T, uint32_t... dims> constant constexpr uint32_t size_bytes = detail::size_info<T, dims...>::bytes;
|
||||
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Functions for transferring data directly between global memory and registers and back.
|
||||
*/
|
||||
#pragma once // not done
|
||||
/*
|
||||
TODO:
|
||||
change loads/stores, prevent unnecessary
|
||||
*/
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
/**
|
||||
* @brief Load data into a register vector 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<typename RV, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>() && ducks::is_global_layout<GL>(), void>::type
|
||||
load(thread RV &dst, thread const GL &src, thread const coord &idx, const short laneid) {
|
||||
using RV_T = typename RV::dtype;
|
||||
using RV_T2 = typename base_types::packing<RV_T>::packed_type;
|
||||
using U = typename GL::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
device U *src_ptr = (device U*)&src.template get<RV>(idx);
|
||||
if (ducks::is_align_layout<typename RV::layout>()) {
|
||||
constexpr const uint32_t MASK_1 = 0x00AA00AA; // kitty bit magic
|
||||
constexpr const uint32_t MASK_2 = 0x55005500;
|
||||
constexpr const uint32_t MASK_3 = 0xAA00AA00;
|
||||
unsigned offset = ((MASK_1 >> laneid) & 1u) * 2 + ((MASK_2 >> laneid) & 1u) * 4 + ((MASK_3 >> laneid) & 1u) * 6;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int t = 0; t < RV::outer_dim; offset+=8, t++) {
|
||||
RV_T2 src2 = base_types::convertor<RV_T2, U2>::convert(*(device U2*)(&src_ptr[offset]));
|
||||
dst.data[t][0] = src2[0];
|
||||
dst.data[t][1] = src2[1];
|
||||
}
|
||||
} else if (ducks::is_ortho_layout<typename RV::layout>()) { // RV::inner_dim == 1
|
||||
const short laneid_div2 = laneid / 2;
|
||||
unsigned offset = laneid_div2 % 4 + (laneid_div2 / 8) * 4;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int t = 0; t < RV::outer_dim; offset+=8, t++) {
|
||||
dst.data[t][0] = base_types::convertor<RV_T, U>::convert(src_ptr[offset]);
|
||||
}
|
||||
} else if (ducks::is_naive_layout<typename RV::layout>()) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(auto w = 0; w < RV::outer_dim; w++) {
|
||||
// if(w < dst.outer_dim-1 || dst.length%32 == 0 || laneid<16) {
|
||||
if (w * SIMD_THREADS + laneid < RV::length) {
|
||||
dst[w][0] = base_types::convertor<RV_T, U>::convert(src_ptr[w * SIMD_THREADS + laneid]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store data from a register vector 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<typename RV, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>() && ducks::is_global_layout<GL>(), void>::type
|
||||
store(thread GL &dst, thread const RV &src, thread const coord &idx, const short laneid) {
|
||||
using RV_T = typename RV::dtype;
|
||||
using RV_T2 = typename base_types::packing<RV_T>::packed_type;
|
||||
using U = typename GL::dtype;
|
||||
using U2 = typename base_types::packing<U>::packed_type;
|
||||
device U *dst_ptr = (device U*)&(dst.template get<RV>(idx));
|
||||
if (ducks::is_align_layout<typename RV::layout>()) {
|
||||
constexpr const uint32_t MASK_1 = 0x00AA00AA; // kitty bit magic
|
||||
constexpr const uint32_t MASK_2 = 0x55005500;
|
||||
constexpr const uint32_t MASK_3 = 0xAA00AA00;
|
||||
unsigned offset = ((MASK_1 >> laneid) & 1u) * 2 + ((MASK_2 >> laneid) & 1u) * 4 + ((MASK_3 >> laneid) & 1u) * 6;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int t = 0; t < RV::outer_dim; offset+=8, t++) {
|
||||
U2 src2 = base_types::convertor<U2, RV_T2>::convert({src.data[t][0], src.data[t][1]});
|
||||
*(device U2*)(&dst_ptr[offset]) = src2;
|
||||
}
|
||||
} else if (ducks::is_ortho_layout<typename RV::layout>()){ // RV::inner_dim == 1
|
||||
const short laneid_div2 = laneid / 2;
|
||||
unsigned offset = laneid_div2 % 4 + (laneid_div2 / 8) * 4;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int t = 0; t < RV::outer_dim; offset+=8, t++) {
|
||||
dst_ptr[offset] = base_types::convertor<U, RV_T>::convert(src.data[t][0]);
|
||||
}
|
||||
} else {
|
||||
#pragma clang loop unroll(full)
|
||||
for(auto w = 0; w < RV::outer_dim; w++) {
|
||||
// if(w < dst.outer_dim-1 || dst.length%32 == 0 || laneid<16) {
|
||||
if (w * SIMD_THREADS + laneid < RV::length) {
|
||||
dst_ptr[w * SIMD_THREADS + laneid] = base_types::convertor<U, RV_T>::convert(src.data[w][0]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Functions for transferring data directly between global and shared memory and back.
|
||||
*/
|
||||
|
||||
#pragma once // done!
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
|
||||
template<typename SV, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_shared_vector<SV>() && ducks::is_global_layout<GL>(), void>::type
|
||||
load(threadgroup SV &dst, thread const GL &src, thread const coord &idx, const unsigned laneid) {
|
||||
using read_type = float4;
|
||||
using U = typename GL::dtype;
|
||||
constexpr int elem_per_transfer = sizeof(read_type) / sizeof(typename SV::dtype);
|
||||
constexpr int total_calls = SV::length / elem_per_transfer; // guaranteed to divide
|
||||
device U *src_ptr = (device U*)&src.template get<SV>(idx);
|
||||
#pragma clang loop unroll(full)
|
||||
for (int i = laneid; i < total_calls; i += mittens::SIMD_THREADS) {
|
||||
if(i * elem_per_transfer < dst.length) {
|
||||
*(threadgroup read_type*)&dst[i*elem_per_transfer] = *(device read_type*)&src_ptr[i*elem_per_transfer];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<typename SV, typename GL>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_shared_vector<SV>() && ducks::is_global_layout<GL>(), void>::type
|
||||
store(thread const GL &dst, threadgroup const SV &src, thread const coord &idx, const unsigned laneid) {
|
||||
using read_type = float4;
|
||||
using U = typename GL::dtype;
|
||||
constexpr int elem_per_transfer = sizeof(read_type) / sizeof(typename SV::dtype);
|
||||
constexpr int total_calls = SV::length / elem_per_transfer; // guaranteed to divide
|
||||
device U *dst_ptr = (device U*)&dst.template get<SV>(idx);
|
||||
#pragma clang loop unroll(full)
|
||||
for (int i = laneid; i < total_calls; i += mittens::SIMD_THREADS) {
|
||||
if(i * elem_per_transfer < src.length) {
|
||||
*(device read_type*)&dst_ptr[i*elem_per_transfer] = *(threadgroup read_type*)&src[i*elem_per_transfer];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Functions for transferring data directly between shared memory and registers and back.
|
||||
*/
|
||||
|
||||
#pragma once // not done
|
||||
/*
|
||||
TODO:
|
||||
prevent unnecesary memory back forth
|
||||
|
||||
*/
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
|
||||
/**
|
||||
* @brief Load data from a shared vector into a register vector.
|
||||
*
|
||||
* @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.
|
||||
*/
|
||||
|
||||
/*
|
||||
"For row-vectors:
|
||||
0,2,4,6,16,18,20,22 holds %8+0 & %8 +1
|
||||
1,3,5,7,17,19,21,23 holds %8+2 & %8+3
|
||||
00000000101010100000000010101010 = 0x00AA00AA
|
||||
8,10,12,14,24,26,28,30 holds %8+4 & %8+5
|
||||
01010101000000000101010100000000 = 0x55005500
|
||||
9,11,13,15,25,27,29,31 holds %8+6 & %8+7"
|
||||
10101010000000001010101000000000 = 0xAA00AA00
|
||||
|
||||
"For colum-vectors:
|
||||
0,1,8,9 holds %8+0
|
||||
2,3,10,11 holds %8+1
|
||||
4,5,12,13 holds %8+2
|
||||
6,7,14,15 holds %8+3
|
||||
16,17,24,25 holds %8+4
|
||||
18,19,26,27 holds %8+5
|
||||
20,21,28,29 holds %8+6
|
||||
22,23,30,31 holds %8+7
|
||||
|
||||
0,0,4,4 holds %8+0
|
||||
1,1,5,5 holds %8+1
|
||||
2,2,6,6 holds %8+2
|
||||
3,3,7,7 holds %8+3
|
||||
8,8,12,12 holds %8+4
|
||||
9,9,13,13 holds %8+5
|
||||
10,10,14,14 holds %8+6
|
||||
11,11,15,15 holds %8+7
|
||||
"
|
||||
|
||||
0 0 1 1 8 8 9 9
|
||||
2 2 3 3 10 10 11 11
|
||||
4 4 5 5 12 12 13 13
|
||||
6 6 7 7 14 14 15 15
|
||||
16 16 17 17 24 24 25 25
|
||||
18 18 19 19 26 26 27 27
|
||||
20 20 21 21 28 28 29 29
|
||||
22 22 23 23 30 30 31 31
|
||||
*/
|
||||
// optimize later
|
||||
template<typename RV, typename SV>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
load(thread RV &dst, threadgroup const SV &src, const short laneid) {
|
||||
using RV_T = typename RV::dtype;
|
||||
using RV_T2 = typename base_types::packing<RV_T>::packed_type;
|
||||
using SV_T = typename SV::dtype;
|
||||
using SV_T2 = typename base_types::packing<SV_T>::packed_type;
|
||||
|
||||
|
||||
static_assert(SV::tiles == RV::tiles, "RV and SV dimensions must match");
|
||||
|
||||
if (ducks::is_align_layout<typename RV::layout>()) {
|
||||
constexpr const uint32_t MASK_1 = 0x00AA00AA; // kitty bit magic
|
||||
constexpr const uint32_t MASK_2 = 0x55005500;
|
||||
constexpr const uint32_t MASK_3 = 0xAA00AA00;
|
||||
unsigned offset = ((MASK_1 >> laneid) & 1u) * 2 + ((MASK_2 >> laneid) & 1u) * 4 + ((MASK_3 >> laneid) & 1u) * 6;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int t = 0; t < SV::tiles; offset+=8, t++) {
|
||||
RV_T2 src2 = base_types::convertor<RV_T2, SV_T2>::convert(*(threadgroup SV_T2*)(&src.data[offset]));
|
||||
dst.data[t][0] = src2[0];
|
||||
dst.data[t][1] = src2[1];
|
||||
// dst.data[t][0] = 7.f;
|
||||
// dst.data[t][1] = 7.f;
|
||||
}
|
||||
} else if (ducks::is_ortho_layout<typename RV::layout>()) {
|
||||
const short laneid_div2 = laneid / 2;
|
||||
unsigned offset = laneid_div2 % 4 + (laneid_div2 / 8) * 4;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int t = 0; t < SV::tiles; offset+=8, t++) {
|
||||
dst.data[t][0] = base_types::convertor<RV_T, SV_T>::convert(src[offset]);
|
||||
}
|
||||
} else if (ducks::is_naive_layout<typename RV::layout>()) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(auto w = 0; w < RV::outer_dim; w++) {
|
||||
if (w * SIMD_THREADS + laneid < RV::length) {
|
||||
dst.data[w][0] = base_types::convertor<RV_T, SV_T>::convert(src[w * SIMD_THREADS + laneid]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* @brief Store data into a shared vector from a register vector.
|
||||
*
|
||||
* @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.
|
||||
*/
|
||||
// optimize later
|
||||
template<typename SV, typename RV>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
store(threadgroup SV &dst, thread const RV &src, const short laneid) {
|
||||
ducks::assert_shared_vector<SV>();
|
||||
ducks::assert_register_vector<RV>();
|
||||
using RV_T = typename RV::dtype;
|
||||
using RV_T2 = typename base_types::packing<RV_T>::packed_type;
|
||||
using SV_T = typename SV::dtype;
|
||||
using SV_T2 = typename base_types::packing<SV_T>::packed_type;
|
||||
|
||||
|
||||
static_assert(SV::tiles == RV::tiles, "RV and SV dimensions must match");
|
||||
|
||||
if (ducks::is_align_layout<typename RV::layout>()) {
|
||||
constexpr const uint32_t MASK_1 = 0x00AA00AA; // kitty bit magic
|
||||
constexpr const uint32_t MASK_2 = 0x55005500;
|
||||
constexpr const uint32_t MASK_3 = 0xAA00AA00;
|
||||
unsigned offset = ((MASK_1 >> laneid) & 1u) * 2 + ((MASK_2 >> laneid) & 1u) * 4 + ((MASK_3 >> laneid) & 1u) * 6;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int t = 0; t < SV::tiles; offset+=8, t++) {
|
||||
SV_T2 src2 = base_types::convertor<SV_T2, RV_T2>::convert({src.data[t][0], src.data[t][1]});
|
||||
*(threadgroup SV_T2*)(&dst.data[offset]) = src2;
|
||||
|
||||
// *(threadgroup SV_T2*)(&dst.data[offset]) = (SV_T2)1.f;
|
||||
}
|
||||
} else if (ducks::is_ortho_layout<typename RV::layout>()) {
|
||||
const short laneid_div2 = laneid / 2;
|
||||
unsigned offset = laneid_div2 % 4 + (laneid_div2 / 8) * 4;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int t = 0; t < SV::tiles; offset+=8, t++) {
|
||||
dst[offset] = base_types::convertor<SV_T, RV_T>::convert(src.data[t][0]);
|
||||
}
|
||||
} else if (ducks::is_naive_layout<typename RV::layout>()) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(auto w = 0; w < RV::outer_dim; w++) {
|
||||
if (w * SIMD_THREADS + laneid < RV::length) {
|
||||
dst[w * SIMD_THREADS + laneid] = base_types::convertor<SV_T, RV_T>::convert(src.data[w][0]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
|
||||
|
||||
///// TRASH CAN
|
||||
|
||||
/*
|
||||
template<typename RV, typename SV>
|
||||
METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
load(thread RV &dst, threadgroup const SV &src, const short laneid, const int start_tile, const int size_tile) {
|
||||
using RV_T = typename RV::dtype;
|
||||
using RV_T2 = typename base_types::packing<RV_T>::packed_type;
|
||||
using SV_T = typename SV::dtype;
|
||||
using SV_T2 = typename base_types::packing<SV_T>::packed_type;
|
||||
|
||||
|
||||
// static_assert(RV::tiles == size_tile , "RV and SV dimensions must match");
|
||||
|
||||
if (ducks::is_align_layout<typename RV::layout>()) {
|
||||
constexpr const uint32_t MASK_1 = 0x00AA00AA; // kitty bit magic
|
||||
constexpr const uint32_t MASK_2 = 0x55005500;
|
||||
constexpr const uint32_t MASK_3 = 0xAA00AA00;
|
||||
unsigned offset = ((MASK_1 >> laneid) & 1u) * 2 + ((MASK_2 >> laneid) & 1u) * 4 + ((MASK_3 >> laneid) & 1u) * 6
|
||||
+ 8 * start_tile;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int t = start_tile; t < start_tile + size_tile; offset+=8, t++) {
|
||||
// RV_T2 src2 = base_types::convertor<RV_T2, SV_T2>::convert(*(threadgroup SV_T2*)(&src.data[offset]));
|
||||
// dst.data[t][0] = src2[0];
|
||||
// dst.data[t][1] = src2[1];
|
||||
}
|
||||
} else if (ducks::is_ortho_layout<typename RV::layout>()) {
|
||||
const short laneid_div2 = laneid / 2;
|
||||
unsigned offset = laneid_div2 % 4 + (laneid_div2 / 8) * 4
|
||||
+ 8 * start_tile;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int t = start_tile; t < start_tile + size_tile; offset+=8, t++) {
|
||||
dst.data[t][0] = base_types::convertor<RV_T, SV_T>::convert(src[offset]);
|
||||
}
|
||||
}
|
||||
// else if (ducks::is_naive_layout<typename RV::layout>()) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(auto w = 0; w < RV::outer_dim; w++) {
|
||||
// if (w * SIMD_THREADS + laneid < RV::length) {
|
||||
// dst.data[w][0] = base_types::convertor<RV_T, SV_T>::convert(src[w * SIMD_THREADS + laneid]);
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
}
|
||||
|
||||
*/
|
||||
@@ -0,0 +1,4 @@
|
||||
#pragma once
|
||||
#include "global_to_register.metal"
|
||||
#include "global_to_shared.metal"
|
||||
#include "shared_to_register.metal"
|
||||
@@ -0,0 +1,3 @@
|
||||
#pragma once
|
||||
#include "tile/tile.metal"
|
||||
#include "vec/vec.metal"
|
||||
@@ -0,0 +1,313 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Conversions between data layouts and types for register tiles.
|
||||
*/
|
||||
|
||||
#pragma once // not done:
|
||||
/*
|
||||
swaping register layout doesn't exist. no layout to swap
|
||||
SUBTILE
|
||||
|
||||
*/
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
/* ---------- TRANSPOSE ---------- */
|
||||
METAL_FUNC int compute_laneid(ushort y, ushort x) {
|
||||
// Extract bits from simd_y
|
||||
ushort b1 = y & 1;
|
||||
ushort temp_y = y >> 1;
|
||||
ushort b2 = temp_y & 1;
|
||||
ushort b4 = temp_y >> 1;
|
||||
|
||||
// Extract bits from simd_x
|
||||
ushort b0 = (x >> 1) & 1;
|
||||
ushort b3 = x >> 2;
|
||||
|
||||
// Reconstruct laneid
|
||||
ushort laneid = (b4 << 4) | (b3 << 3) | (b2 << 2) | (b1 << 1) | b0;
|
||||
return laneid;
|
||||
}
|
||||
/**
|
||||
* @brief Transposes a register base tile.
|
||||
*
|
||||
* @tparam T2 The data type of the register tile elements.
|
||||
* @tparam layout The current layout of the register tile.
|
||||
* @param dst[out] Reference to the register tile in which to store the transposed src.
|
||||
* @param src[in] Reference to the register base tile to be transposed.
|
||||
*/
|
||||
template<typename T, typename layout>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_rt_layout<layout>(), void>::type
|
||||
swap_layout(thread rt_base<T, typename ducks::rt_layout::transpose<layout>::type> &dst,
|
||||
thread const rt_base<T, layout> &src,
|
||||
const ushort laneid) {
|
||||
const ushort qid = laneid / 4;
|
||||
const ushort simd_y = (qid & 4) + (laneid / 2) % 4;
|
||||
const ushort simd_x = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
|
||||
const ushort src_laneid_start = compute_laneid(simd_x, simd_y);
|
||||
const ushort2 src_laneid = ushort2(src_laneid_start, src_laneid_start+(ushort)2);
|
||||
const ushort first_idx = (laneid / 2) % 2;
|
||||
|
||||
dst.data.thread_elements()[first_idx] = shfl_sync<T>(src.data.thread_elements()[first_idx], src_laneid[first_idx]);
|
||||
|
||||
dst.data.thread_elements()[1 - first_idx] = shfl_sync<T>(src.data.thread_elements()[1 - first_idx], src_laneid[1 - first_idx]);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Swaps the layout of a register tile.
|
||||
*
|
||||
* This function swaps the layout of a register tile by iterating over its height and width
|
||||
* and performing layout swaps on each of its base elements.
|
||||
*
|
||||
* @tparam T2 The data type of the register tile elements.
|
||||
* @tparam _height The height of the register tile.
|
||||
* @tparam _width The width of the register tile.
|
||||
* @tparam layout The current layout of the register tile.
|
||||
* @param dst[out] Reference to the destination register tile where the result will be stored.
|
||||
* @param src[in] Reference to the source register tile to be swapped.
|
||||
*/
|
||||
template<typename T, int _height, int _width, typename layout>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_rt_layout<layout>(), void>::type
|
||||
swap_layout(thread rt<T, _height, _width, typename ducks::rt_layout::transpose<layout>::type> &dst, thread const rt<T, _height, _width, layout> &src, const short laneid) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
swap_layout(dst.tiles[i][j], src.tiles[i][j], laneid);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Swaps the layout of a register base tile in place.
|
||||
*
|
||||
* This function swaps the layout of a register base tile in place by casting it to the
|
||||
* transposed layout type and then performing the layout swap.
|
||||
*
|
||||
* @tparam T2 The data type of the register tile elements.
|
||||
* @tparam layout The current layout of the register tile.
|
||||
* @param src[in] Reference to the register base tile to be swapped in place.
|
||||
* @return A reference to the swapped register base tile.
|
||||
*/
|
||||
template<typename T2, typename layout>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_rt_layout<layout>(), thread rt_base<T2, typename ducks::rt_layout::transpose<layout>::type>&>::type
|
||||
swap_layout_inplace(thread const rt_base<T2, layout> &src) {
|
||||
thread rt_base<T2, typename ducks::rt_layout::transpose<layout>::type> &dst = *(thread rt_base<T2, typename ducks::rt_layout::transpose<layout>::type>*)(&src);
|
||||
swap_layout(dst, src);
|
||||
return dst;
|
||||
}
|
||||
|
||||
/* ---------- TRANSPOSE ---------- */
|
||||
|
||||
/**
|
||||
* @brief Transposes a register base tile.
|
||||
*
|
||||
* @tparam T2 The data type of the register tile elements.
|
||||
* @tparam layout The current layout of the register tile.
|
||||
* @param dst[out] Reference to the register tile in which to store the transposed src.
|
||||
* @param src[in] Reference to the register base tile to be transposed.
|
||||
*/
|
||||
template<typename T, typename layout>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_rt_layout<layout>(), void>::type
|
||||
transpose(thread rt_base<T, layout> &dst, thread const rt_base<T, layout> &src, const ushort laneid) {
|
||||
const ushort qid = laneid / 4;
|
||||
const ushort simd_y = (qid & 4) + (laneid / 2) % 4;
|
||||
const ushort simd_x = (qid & 2) * 2 + (laneid % 2) * 2;
|
||||
|
||||
const ushort src_laneid_start = compute_laneid(simd_x, simd_y);
|
||||
const ushort2 src_laneid = ushort2(src_laneid_start, src_laneid_start+(ushort)2);
|
||||
const ushort first_idx = (laneid / 2) % 2;
|
||||
|
||||
dst.data.thread_elements()[first_idx] = shfl_sync<T>(src.data.thread_elements()[first_idx], src_laneid[first_idx]);
|
||||
|
||||
dst.data.thread_elements()[1 - first_idx] = shfl_sync<T>(src.data.thread_elements()[1 - first_idx], src_laneid[1 - first_idx]);
|
||||
}
|
||||
/**
|
||||
* @brief Transposes a register tile.
|
||||
*
|
||||
* This function is marked "sep", which means that the registers underlying dst MUST be separate
|
||||
* from the registers underlying src.
|
||||
*
|
||||
* @tparam T2 The data type of the register tile elements.
|
||||
* @tparam _height The height of the src register tile, and the width of the dst tile.
|
||||
* @tparam _width The width of the src register tile, and the height of the dst tile.
|
||||
* @tparam layout The layout of the register tile.
|
||||
* @param dst[out] Reference to the register tile in which to store the transposed src.
|
||||
* @param src[in] Reference to the register tile to be transposed.
|
||||
*/
|
||||
template<typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
transpose_sep(thread RT &dst, thread const rt<typename RT::T, RT::cols, RT::rows, typename RT::layout> &src,
|
||||
const int laneid) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < RT::height; i++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < RT::width; j++) {
|
||||
transpose(dst.tiles[i][j], src.tiles[j][i], laneid);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Transposes a register base tile in-place.
|
||||
*
|
||||
* @tparam T2 The data type of the register base tile elements.
|
||||
* @tparam layout The current layout of the register base tile.
|
||||
* @param src[in] Reference to the register tile to be transposed.
|
||||
* @return A reference to the transposed register base tile.
|
||||
*/
|
||||
template<typename T2, typename layout>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_rt_layout<layout>(), thread rt_base<T2, layout>&>::type
|
||||
transpose_inplace(thread rt_base<T2, layout> &src, const ushort laneid) {
|
||||
transpose(src, src, laneid);
|
||||
return src;
|
||||
}
|
||||
|
||||
template<typename T, typename U, typename layout>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_rt_layout<layout>(), void>::type
|
||||
copy(thread rt_base<T, layout> &dst, thread const rt_base<U, layout> &src);
|
||||
|
||||
/**
|
||||
* @brief Transposes a square register tile in-place.
|
||||
*
|
||||
* @tparam T2 The data type of the register tile elements.
|
||||
* @tparam _height The height (in units of 16) of the src register tile, and the width of the dst tile. (Must be the same as _width.)
|
||||
* @tparam _width The width (in units of 16) of the src register tile, and the height of the dst tile. (Must be the same as _height.)
|
||||
* @tparam layout The current layout of the register tile.
|
||||
* @param src[in] Reference to the register tile to be transposed.
|
||||
* @return A reference to the transposed register tile.
|
||||
*/
|
||||
template<typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && RT::cols == RT::rows, thread RT&>::type
|
||||
transpose_inplace(thread RT &tile, const ushort laneid) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < tile.height; i++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < i; j++) {
|
||||
rt_base<typename RT::T, typename RT::layout> tmp;
|
||||
copy(tmp, tile.tiles[i][j]);
|
||||
transpose(tile.tiles[i][j], tile.tiles[j][i], laneid);
|
||||
transpose(tile.tiles[j][i], tmp, laneid);
|
||||
}
|
||||
transpose_inplace(tile.tiles[i][i], laneid);
|
||||
}
|
||||
return tile;
|
||||
}
|
||||
/* ---------- TYPE SWAPS ---------- */
|
||||
/**
|
||||
* @brief Copies a register base tile, converting the underlying type if necessary.
|
||||
*
|
||||
* @tparam T2 The data type of the destination register elements.
|
||||
* @tparam U2 The data type of the source register elements.
|
||||
* @tparam layout The current layout of the register base tile.
|
||||
* @param[out] dst A reference to the destination register base tile.
|
||||
* @param[in] src A reference to the source register base tile.
|
||||
*/
|
||||
template<typename T, typename U, typename layout>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_rt_layout<layout>(), void>::type
|
||||
copy(thread rt_base<T, layout> &dst, thread const rt_base<U, layout> &src) {
|
||||
using T1 = typename base_types::packing<T>::unpacked_type;
|
||||
using U1 = typename base_types::packing<U>::unpacked_type;
|
||||
dst.data.thread_elements()[0] = base_types::convertor<T1, U1>::convert(src.data.thread_elements()[0]);
|
||||
dst.data.thread_elements()[1] = base_types::convertor<T1, U1>::convert(src.data.thread_elements()[1]);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Copies a register tile, converting the underlying type if necessary.
|
||||
*
|
||||
* @tparam T2 The data type of the destination register elements.
|
||||
* @tparam U2 The data type of the source register elements.
|
||||
* @tparam _height The height (in units of 8) of the register tiles.
|
||||
* @tparam _width The width (in units of 8) of the register tiles.
|
||||
* @tparam layout The current layout of the register tile.
|
||||
* @param[out] dst A reference to the destination register tile.
|
||||
* @param[in] src A reference to the source register tile.
|
||||
*/
|
||||
template<typename T, typename U, int _height, int _width, typename layout>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_rt_layout<layout>(), void>::type
|
||||
copy(thread rt<T, _height, _width, layout> &dst, thread const rt<U, _height, _width, layout> &src) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
copy(dst.tiles[i][j], src.tiles[i][j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------- CAUSAL ---------- */
|
||||
|
||||
/**
|
||||
* @brief Makes a square register tile causal by zeroing elements above the main diagonal.
|
||||
*
|
||||
* This function modifies a square register tile in-place to make it causal. All elements
|
||||
* above the main diagonal are set to zero, while elements on or below the main diagonal
|
||||
* are left unchanged.
|
||||
*
|
||||
* @tparam T The data type of the register tile elements.
|
||||
* @tparam _size The size (height and width) of the square register tile.
|
||||
* @tparam layout The current layout of the register tile.
|
||||
* @param tile[in,out] Reference to the register tile to be made causal.
|
||||
*/
|
||||
template<typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
make_causal(thread RT &dst, thread const RT &src, const unsigned laneid, thread const typename base_types::packing<typename RT::dtype>::unpacked_type &val=0) {
|
||||
ducks::assert_register_tile<RT>();
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
if(j < i) { // below the diagonal, copy
|
||||
dst.tiles[i][j].data.thread_elements()[0] = src.tiles[i][j].data.thread_elements()[0];
|
||||
dst.tiles[i][j].data.thread_elements()[1] = src.tiles[i][j].data.thread_elements()[1];
|
||||
}
|
||||
else if(j > i) { // above the diagonal, zero
|
||||
dst.tiles[i][j].data.thread_elements()[0] = val;
|
||||
dst.tiles[i][j].data.thread_elements()[1] = val;
|
||||
}
|
||||
else { // on the diagonal
|
||||
constexpr uint32_t MASK_0 = (ducks::is_row_register_tile<RT>()) ? 0x0A00FF0A : 0xD4FF00D4;
|
||||
constexpr uint32_t MASK_1 = (ducks::is_row_register_tile<RT>()) ? 0x2B00FF2B : 0x50FF0050;
|
||||
if((MASK_0 >> laneid) & 1) {
|
||||
dst.tiles[i][j].data.thread_elements()[0] = val;
|
||||
}
|
||||
else {
|
||||
dst.tiles[i][j].data.thread_elements()[0] = src.tiles[i][j].data.thread_elements()[0];
|
||||
}
|
||||
if((MASK_1 >> laneid) & 1) {
|
||||
dst.tiles[i][j].data.thread_elements()[1] = val;
|
||||
}
|
||||
else {
|
||||
dst.tiles[i][j].data.thread_elements()[1] = src.tiles[i][j].data.thread_elements()[1];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/* ---------- SUBTILE ---------- */
|
||||
|
||||
/**
|
||||
* @brief Returns a reference to a subtile of the given tile.
|
||||
*
|
||||
* @tparam subtile_height The height of the subtile.
|
||||
* @tparam RT The type of the input tile, which must satisfy the ducks::rt::all concept.
|
||||
* @param src The input tile.
|
||||
* @param idx The index of the subtile.
|
||||
* @return A reference to the subtile.
|
||||
*
|
||||
* @note The subtile height must evenly divide the tile height.
|
||||
*/
|
||||
//template<int subtile_height, ducks::rt::all RT>
|
||||
//__device__ inline rt<typename RT::T, subtile_height, RT::width, typename RT::layout> &subtile_inplace(RT & src, int idx) {
|
||||
// static_assert(RT::height % subtile_height == 0, "subtile height should evenly divide tile height.");
|
||||
// return reinterpret_cast<rt<typename RT::T, subtile_height, RT::width, typename RT::layout>&>(
|
||||
// src.tiles[idx*subtile_height]
|
||||
// );
|
||||
//}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,878 @@
|
||||
#pragma once // doneington but add register tile col
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
/* ---------- Uniform tile maps (independent of layout) ---------- */
|
||||
|
||||
namespace meta {
|
||||
template<typename op, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
unary_map_unroll(int i, int j, thread RT *dst, thread const RT *src) {
|
||||
using T2 = typename RT::T2;
|
||||
T2 vals = op::template op<T2>(T2{src->tiles[i][j].data.thread_elements()[0], src->tiles[i][j].data.thread_elements()[1]});
|
||||
dst->tiles[i][j].data.thread_elements()[0] = vals[0];
|
||||
dst->tiles[i][j].data.thread_elements()[1] = vals[1];
|
||||
}
|
||||
}
|
||||
/**
|
||||
* @brief Applies a unary operation to each element of a tile.
|
||||
*
|
||||
* @tparam op Unary operation to apply.
|
||||
* @tparam T Tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the operation on.
|
||||
*/
|
||||
template<typename op, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
unary_map(thread RT &dst, thread const RT &src) {
|
||||
using T = typename RT::T;
|
||||
ducks::assert_register_tile<RT>();
|
||||
using T2 = typename RT::T2;
|
||||
using T4 = typename base_types::packing<typename RT::dtype>::packed_four;
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < dst.height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < dst.width; j++) {
|
||||
// T2 op2 = op::template op<T2>(T2{src.tiles[i][j].data.thread_elements()[0], src.tiles[i][j].data.thread_elements()[1]});
|
||||
//// dst.tiles[i][j].data.thread_elements()[0] = op::template op<typename RT::dtype>(src.tiles[i][j].data.thread_elements()[0]);
|
||||
//// dst.tiles[i][j].data.thread_elements()[1] = op::template op<typename RT::dtype>(src.tiles[i][j].data.thread_elements()[1]);
|
||||
//
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = op2[0];
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = op2[1];
|
||||
//
|
||||
//// dst.tiles[i][j].data.thread_elements()[0] = base_ops::abs::template op<T>(src.tiles[i][j].data.thread_elements()[0]);
|
||||
//// dst.tiles[i][j].data.thread_elements()[1] = base_ops::abs::template op<T>(src.tiles[i][j].data.thread_elements()[1]);
|
||||
//// dst.tiles[i][j].data.thread_elements()[0] = (T)(metal::abs(-1.f));
|
||||
//// dst.tiles[i][j].data.thread_elements()[1] = (T)(metal::abs(-1.f));
|
||||
//
|
||||
//// ((T)(((float)src.tiles[i][j].data.thread_elements()[0])));
|
||||
//// dst.tiles[i][j].data.thread_elements()[1] = metal::abs((T)((float)src.tiles[i][j].data.thread_elements()[1]));
|
||||
//
|
||||
//// dst.tiles[i][j].data.thread_elements()[0] = base_types::constants<typename RT::dtype>::one();
|
||||
//// dst.tiles[i][j].data.thread_elements()[1] = base_types::constants<typename RT::dtype>::one();
|
||||
//// metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
//
|
||||
////// T2 val = op::template op<T2>(T2{src.tiles[i][j].data.thread_elements()[0],
|
||||
////// src.tiles[i][j].data.thread_elements()[1]});
|
||||
////// dst.tiles[i][j].data.thread_elements()[0] = val[0];
|
||||
////// dst.tiles[i][j].data.thread_elements()[1] = val[1];
|
||||
////////
|
||||
////// T4 val = op::template op<T4>(T4{src.tiles[i][j].data.thread_elements()[0],
|
||||
////// src.tiles[i][j].data.thread_elements()[1],
|
||||
////// src.tiles[i][j+1].data.thread_elements()[0],
|
||||
////// src.tiles[i][j+1].data.thread_elements()[1],});
|
||||
////// dst.tiles[i][j].data.thread_elements()[0] = val[0];
|
||||
////// dst.tiles[i][j].data.thread_elements()[1] = val[1];
|
||||
////// dst.tiles[i][j+1].data.thread_elements()[0] = val[2];
|
||||
////// dst.tiles[i][j+1].data.thread_elements()[1] = val[3];
|
||||
// }
|
||||
// }
|
||||
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::unary_map_unroll<op, RT>, &dst, &src);
|
||||
}
|
||||
|
||||
|
||||
namespace meta {
|
||||
template<typename op, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
bin_map_unroll(int i, int j, thread RT *dst, thread const RT *src, thread const typename RT::dtype *param) {
|
||||
using T = typename RT::T;
|
||||
using T2 = typename RT::T2;
|
||||
// T2 vals = op::template op<T2>({src->tiles[i][j].data.thread_elements()[0], src->tiles[i][j].data.thread_elements()[1]}, {*param, *param});
|
||||
// dst->tiles[i][j].data.thread_elements()[0] = vals[0];
|
||||
// dst->tiles[i][j].data.thread_elements()[1] = vals[1];
|
||||
dst->tiles[i][j].data.thread_elements()[0] = op::template op<T>(src->tiles[i][j].data.thread_elements()[0], *param);
|
||||
dst->tiles[i][j].data.thread_elements()[1] = op::template op<T>(src->tiles[i][j].data.thread_elements()[1], *param);
|
||||
}
|
||||
}
|
||||
/**
|
||||
* @brief Applies a binary operation to each element of a tile with a scalar parameter.
|
||||
*
|
||||
* @tparam op Binary operation to apply.
|
||||
* @tparam T Tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the operation on.
|
||||
* @param param[in] Scalar parameter for the binary operation.
|
||||
*/
|
||||
template<typename op, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
bin_map(thread RT &dst, thread const RT &src, thread const typename RT::dtype ¶m) {
|
||||
// using T = typename RT::T;
|
||||
// using T2 = typename RT::T2;
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < dst.height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < dst.width; j++) {
|
||||
// T2 vals = op::template op<T2>({src.tiles[i][j].data.thread_elements()[0], src.tiles[i][j].data.thread_elements()[1]}, {param, param});
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = vals[0];
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = vals[1];
|
||||
//// dst.tiles[i][j].data.thread_elements()[0] = op::template op<typename RT::dtype>(src.tiles[i][j].data.thread_elements()[0], param);
|
||||
//// dst.tiles[i][j].data.thread_elements()[1] = op::template op<typename RT::dtype>(src.tiles[i][j].data.thread_elements()[1], param);
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::bin_map_unroll<op, RT>, &dst, &src, ¶m);
|
||||
}
|
||||
|
||||
namespace meta {
|
||||
template<typename op, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
binary_map_unroll(int i, int j, thread RT *dst, thread const RT *lhs, thread const RT *rhs) {
|
||||
using T2 = typename RT::T2;
|
||||
using T4 = typename base_types::packing<typename RT::dtype>::packed_four;
|
||||
dst->tiles[i][j].data.thread_elements()[0] = op::template op<typename RT::dtype>(lhs->tiles[i][j].data.thread_elements()[0],
|
||||
rhs->tiles[i][j].data.thread_elements()[0]);
|
||||
dst->tiles[i][j].data.thread_elements()[1] = op::template op<typename RT::dtype>(lhs->tiles[i][j].data.thread_elements()[1],
|
||||
rhs->tiles[i][j].data.thread_elements()[1]);
|
||||
// T2 vals = op::template op<T2>({lhs->tiles[i][j].data.thread_elements()[0], lhs->tiles[i][j].data.thread_elements()[1]},
|
||||
// {rhs->tiles[i][j].data.thread_elements()[0], rhs->tiles[i][j].data.thread_elements()[1]});
|
||||
////
|
||||
// dst->tiles[i][j].data.thread_elements()[0] = vals[0];
|
||||
// dst->tiles[i][j].data.thread_elements()[1] = vals[1];
|
||||
|
||||
// dst->tiles[i][j].data.thread_elements()[0] = op::template op<typename RT::dtype>(lhs->tiles[i][j].data.thread_elements()[0],
|
||||
// rhs->tiles[i][j].data.thread_elements()[0]);
|
||||
// dst->tiles[i][j].data.thread_elements()[1] = op::template op<typename RT::dtype>(lhs->tiles[i][j].data.thread_elements()[1],
|
||||
// rhs->tiles[i][j].data.thread_elements()[1]);
|
||||
// T4 val = op::template op<T4>(T4{src->tiles[i][j].data.thread_elements()[0],
|
||||
// src->tiles[i][j].data.thread_elements()[1],
|
||||
// src->tiles[i][j+1].data.thread_elements()[0],
|
||||
// src->tiles[i][j+1].data.thread_elements()[1]});
|
||||
// dst->tiles[i][j].data.thread_elements()[0] = val[0];
|
||||
// dst->tiles[i][j].data.thread_elements()[1] = val[1];
|
||||
// dst->tiles[i][j+1].data.thread_elements()[0] = val[2];
|
||||
// dst->tiles[i][j+1].data.thread_elements()[1] = val[3];
|
||||
}
|
||||
}
|
||||
/**
|
||||
* @brief Applies a binary operation element-wise between two tiles.
|
||||
*
|
||||
* @tparam op Binary operation to apply.
|
||||
* @tparam T Tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param lhs[in] Left-hand side source tile for the operation.
|
||||
* @param rhs[in] Right-hand side source tile for the operation.
|
||||
*/
|
||||
template<typename op, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
bin_map(thread RT &dst, thread const RT &lhs, thread const RT &rhs) {
|
||||
using T = typename RT::dtype;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < dst.height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < dst.width; j++) {
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = op::template op<typename RT::dtype>(lhs.tiles[i][j].data.thread_elements()[0],
|
||||
// rhs.tiles[i][j].data.thread_elements()[0]);
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = op::template op<typename RT::dtype>(lhs.tiles[i][j].data.thread_elements()[1],
|
||||
// rhs.tiles[i][j].data.thread_elements()[1]);
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = lhs.tiles[i][j].data.thread_elements()[0] + rhs.tiles[i][j].data.thread_elements()[0];
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = lhs.tiles[i][j].data.thread_elements()[1] + rhs.tiles[i][j].data.thread_elements()[1];
|
||||
////
|
||||
// T2 vals = op::template op<T2>(T2(lhs.tiles[i][j].data.thread_elements()[0], lhs.tiles[i][j].data.thread_elements()[1]),
|
||||
// T2(rhs.tiles[i][j].data.thread_elements()[0], rhs.tiles[i][j].data.thread_elements()[1]));
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = vals[0];
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = vals[1];
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::binary_map_unroll<op, RT>, &dst, &lhs, &rhs);
|
||||
}
|
||||
|
||||
/* ---------- Row tile maps ----------*/
|
||||
|
||||
namespace meta {
|
||||
template<typename op, typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_map_unroll(int i, int j, thread RT *dst, thread const RT *src, thread const RV *row_values) {
|
||||
using T2 = typename RT::T2;
|
||||
T2 val = op::template op<T2>({src->tiles[i][j].data.thread_elements()[0], src->tiles[i][j].data.thread_elements()[1]}, {(*row_values)[i][0], (*row_values)[i][0]});
|
||||
dst->tiles[i][j].data.thread_elements()[0] = val[0];
|
||||
dst->tiles[i][j].data.thread_elements()[1] = val[1];
|
||||
}
|
||||
|
||||
}
|
||||
/**
|
||||
* @brief Applies an operation across the rows of a tile in a row-major layout.
|
||||
*
|
||||
* @tparam op Operation to apply.
|
||||
* @tparam T Tile type with row-major layout.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the operation on.
|
||||
* @param row_values[in] Column vector containing values to apply across each row.
|
||||
*/
|
||||
template<typename op, typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_map(thread RT &dst, thread const RT &src, thread const RV &row_values) {
|
||||
static_assert(ducks::is_ortho_layout<typename RV::layout>(), "RV must be otho layout (col vec for row rt)");
|
||||
static_assert(metal::is_same_v<typename RV::dtype, typename RT::dtype>, "rt and rv must be of same type"); // compatible type
|
||||
static_assert(RV::outer_dim == RT::height, "RV outer dim and RT height do not match"); // compatible size
|
||||
using T4 = typename base_types::packing<typename RT::dtype>::packed_four;
|
||||
using T2 = typename RT::T2;
|
||||
using T = typename RT::dtype;
|
||||
|
||||
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < RT::height; i++) {
|
||||
// T row_val = row_values[i][0];
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < RT::width; j++) {
|
||||
// T2 val = op::template op<T2>({src.tiles[i][j].data.thread_elements()[0], src.tiles[i][j].data.thread_elements()[1]}, {row_val, row_val});
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = val[0];
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = val[1];
|
||||
//// dst.tiles[i][j].data.thread_elements()[0] = op::template op<T>(src.tiles[i][j].data.thread_elements()[0], row_values[i][0]);
|
||||
//// dst.tiles[i][j].data.thread_elements()[1] = op::template op<T>(src.tiles[i][j].data.thread_elements()[1], row_values[i][0]);
|
||||
// }
|
||||
// }
|
||||
meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::row_map_unroll<op, RT, RV>, &dst, &src, &row_values);
|
||||
|
||||
|
||||
// meta::unroll_i_j_in_range<0, RT::height, 1,
|
||||
// 0, (RT::width / 2) * 2, 2>::run(meta::row_map_unroll<op, RT, RV, 0, 1>, &dst, &src, &row_values);
|
||||
// meta::unroll_i_j_in_range<0, (RT::height / 2) * 2, 2,
|
||||
// (RT::width / 2) * 2, RT::width, 1>::run(meta::row_map_unroll<op, RT, RV, 1, 0>, &dst, &src, &row_values);
|
||||
//
|
||||
// meta::unroll_i_j_in_range<(RT::height / 2) * 2, RT::height, 1,
|
||||
// (RT::width / 2) * 2, RT::width, 1>::run(meta::row_map_unroll<op, RT, RV>, &dst, &src, &row_values);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Applies an operation across the rows of a tile in a row-major layout.
|
||||
*
|
||||
* @tparam op Operation to apply.
|
||||
* @tparam T Tile type with row-major layout.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the operation on.
|
||||
* @param row_values[in] Column vector containing values to apply across each row.
|
||||
*/
|
||||
template<typename op, typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_map(thread RT &dst, thread const RT &src, thread const RV &row_values) {
|
||||
static_assert(ducks::is_align_layout<typename RV::layout>(), "RV must be align layout (col vec for col rt)");
|
||||
static_assert(metal::is_same_v<typename RV::dtype, typename RT::dtype>, "rt and rv must be of same type"); // compatible type
|
||||
static_assert(RV::outer_dim == RT::height, "RV outer dim and RT height do not match"); // compatible size
|
||||
using T4 = typename base_types::packing<typename RT::dtype>::packed_four;
|
||||
using T2 = typename RT::T2;
|
||||
using T = typename RT::dtype;
|
||||
|
||||
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < RT::height; i++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < RT::width; j++) {
|
||||
dst.tiles[i][j].data.thread_elements()[0] = op::template op<T>(src.tiles[i][j].data.thread_elements()[0], row_values[i][0]);
|
||||
dst.tiles[i][j].data.thread_elements()[1] = op::template op<T>(src.tiles[i][j].data.thread_elements()[1], row_values[i][1]);
|
||||
}
|
||||
}
|
||||
//
|
||||
// meta::unroll_i_j_in_range<0, RT::height, 1,
|
||||
// 0, (RT::width / 2) * 2, 2>::run(meta::row_map_unroll<op, RT, RV, 0, 1>, &dst, &src, &row_values);
|
||||
// meta::unroll_i_j_in_range<0, (RT::height / 2) * 2, 2,
|
||||
// (RT::width / 2) * 2, RT::width, 1>::run(meta::row_map_unroll<op, RT, RV, 1, 0>, &dst, &src, &row_values);
|
||||
//
|
||||
// meta::unroll_i_j_in_range<(RT::height / 2) * 2, RT::height, 1,
|
||||
// (RT::width / 2) * 2, RT::width, 1>::run(meta::row_map_unroll<op, RT, RV>, &dst, &src, &row_values);
|
||||
}
|
||||
|
||||
// Three-operand row map. Mostly useful for FMA instructions.
|
||||
|
||||
/**
|
||||
* @brief Applies an operation across the rows of two tiles in a row-major layout, using a third operand.
|
||||
*
|
||||
* @tparam op Operation to apply.
|
||||
* @tparam T Tile type with row-major layout.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param a[in] First source tile to apply the operation on.
|
||||
* @param b[in] Second source tile to apply the operation on.
|
||||
* @param row_values[in] Column vector containing values to apply across each row.
|
||||
*/
|
||||
template<typename op, typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_map(thread RT &dst, thread const RT &a, thread const RT &b, thread const RV &row_values) {
|
||||
static_assert(ducks::is_ortho_layout<RV::layout>(), "rv must be ortho layout for row rt");
|
||||
static_assert(metal::is_same_v<typename RV::dtype, typename RT::dtype>, "rt and rv must be same type"); // compatible type
|
||||
static_assert(RV::outer_dim == RT::height, "rv and rt dimensions don't match"); // compatible size
|
||||
|
||||
|
||||
using dtype = typename RT::dtype;
|
||||
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
dtype vec_val = row_values[i][0];
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
dst.tiles[i][j].data.thread_elements()[0] = op::template op<dtype>(a.tiles[i][j].data.thread_elements()[0], b.tiles[i][j].data.thread_elements()[0], vec_val);
|
||||
|
||||
dst.tiles[i][j].data.thread_elements()[1] = op::template op<dtype>(a.tiles[i][j].data.thread_elements()[1], b.tiles[i][j].data.thread_elements()[1], vec_val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Applies an operation across the rows of two tiles in a column-major layout, using a third operand.
|
||||
*
|
||||
* @tparam op Operation to apply.
|
||||
* @tparam T Tile type with column-major layout.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param a[in] First source tile to apply the operation on.
|
||||
* @param b[in] Second source tile to apply the operation on.
|
||||
* @param row_values[in] Column vector containing values to apply across each row.
|
||||
*/
|
||||
template<typename op, typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_map(thread RT &dst, thread const RT &a, thread const RT &b, thread const RV &row_values) {
|
||||
static_assert(ducks::is_align_layout<RV::layout>(), "rv must be align layout for row rt");
|
||||
static_assert(metal::is_same_v<typename RV::dtype, typename RT::dtype>, "rt and rv must be same type"); // compatible type
|
||||
static_assert(RV::outer_dim == RT::height, "rv and rt dimensions don't match"); // compatible size
|
||||
|
||||
|
||||
using dtype = typename RT::dtype;
|
||||
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
dst.tiles[i][j].data.thread_elements()[0] = op::template op<dtype>(a.tiles[i][j].data.thread_elements()[0], b.tiles[i][j].data.thread_elements()[0], row_values[i][0]);
|
||||
|
||||
dst.tiles[i][j].data.thread_elements()[1] = op::template op<dtype>(a.tiles[i][j].data.thread_elements()[1], b.tiles[i][j].data.thread_elements()[1], row_values[i][1]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------- Col major tile maps ----------*/
|
||||
|
||||
/**
|
||||
* @brief Applies an operation across the columns of a tile in a row-major layout.
|
||||
*
|
||||
* @tparam op Operation to apply.
|
||||
* @tparam T Tile type with row-major layout.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the operation on.
|
||||
* @param col_values[in] Row vector containing values to apply across each column.
|
||||
*/
|
||||
template<typename op, typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_map(thread RT &dst, thread const RT &src, thread const RV &col_values) {
|
||||
static_assert(ducks::is_align_layout<typename RV::layout>(), "rv must be align layout for row rt"); // compatible type
|
||||
static_assert(metal::is_same_v<typename RV::dtype, typename RT::dtype>, "rv and rt must be of the same type"); // compatible type
|
||||
static_assert(RV::outer_dim == RT::width, "rv and rt dimensions do not match"); // compatible size
|
||||
|
||||
using dtype = typename RT::dtype;
|
||||
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
dst.tiles[i][j].data.thread_elements()[0] = op::template op<dtype>(src.tiles[i][j].data.thread_elements()[0], col_values[j][0]);
|
||||
dst.tiles[i][j].data.thread_elements()[1] = op::template op<dtype>(src.tiles[i][j].data.thread_elements()[1], col_values[j][1]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Applies an operation across the columns of a tile in a col-major layout.
|
||||
*
|
||||
* @tparam op Operation to apply.
|
||||
* @tparam T Tile type with row-major layout.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the operation on.
|
||||
* @param col_values[in] Row vector containing values to apply across each column.
|
||||
*/
|
||||
template<typename op, typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_map(thread RT &dst, thread const RT &src, thread const RV &col_values) {
|
||||
static_assert(ducks::is_ortho_layout<typename RV::layout>(), "rv must be ortho layout for row rt"); // compatible type
|
||||
static_assert(metal::is_same_v<typename RV::dtype, typename RT::dtype>, "rv and rt must be of the same type"); // compatible type
|
||||
static_assert(RV::outer_dim == RT::width, "rv and rt dimensions do not match"); // compatible size
|
||||
|
||||
using dtype = typename RT::dtype;
|
||||
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
dst.tiles[i][j].data.thread_elements()[0] = op::template op<dtype>(src.tiles[i][j].data.thread_elements()[0], col_values[j][0]);
|
||||
dst.tiles[i][j].data.thread_elements()[1] = op::template op<dtype>(src.tiles[i][j].data.thread_elements()[1], col_values[j][0]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Three-operand col map
|
||||
/**
|
||||
* @brief Applies an operation across the columns of two tiles in a row-major layout, using a third operand.
|
||||
*
|
||||
* @tparam op Operation to apply.
|
||||
* @tparam T Tile type with row-major layout.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param a[in] First source tile to apply the operation on.
|
||||
* @param b[in] Second source tile to apply the operation on.
|
||||
* @param col_values[in] Row vector containing values to apply across each column.
|
||||
*/
|
||||
template<typename op, typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_map(thread RT &dst, thread const RT &a, thread const RT &b, thread const RV &col_values) {
|
||||
static_assert(ducks::is_align_layout<RV::layout>(), "rv must be align layout");
|
||||
static_assert(metal::is_same_v<typename RV::dtype, typename RT::dtype>, "rv and rt must be of the same type"); // compatible type
|
||||
static_assert(RV::outer_dim == RT::width, "rv and rt dims don't match"); // compatible size
|
||||
|
||||
|
||||
using dtype = typename RT::dtype;
|
||||
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
dst.tiles[i][j].data.thread_elements()[0] = op::template op<dtype>(a.tiles[i][j].data.thread_elements()[0], b.tiles[i][j].data.thread_elements()[0], col_values[j][0]);
|
||||
dst.tiles[i][j].data.thread_elements()[1] = op::template op<dtype>(a.tiles[i][j].data.thread_elements()[1], b.tiles[i][j].data.thread_elements()[1], col_values[j][1]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Applies an operation across the columns of two tiles in a row-major layout, using a third operand.
|
||||
*
|
||||
* @tparam op Operation to apply.
|
||||
* @tparam T Tile type with row-major layout.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param a[in] First source tile to apply the operation on.
|
||||
* @param b[in] Second source tile to apply the operation on.
|
||||
* @param col_values[in] Row vector containing values to apply across each column.
|
||||
*/
|
||||
template<typename op, typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_map(thread RT &dst, thread const RT &a, thread const RT &b, thread const RV &col_values) {
|
||||
static_assert(ducks::is_ortho_layout<RV::layout>(), "rv must be ortho layout");
|
||||
static_assert(metal::is_same_v<typename RV::dtype, typename RT::dtype>, "rv and rt must be of the same type"); // compatible type
|
||||
static_assert(RV::outer_dim == RT::width, "rv and rt dims don't match"); // compatible size
|
||||
|
||||
|
||||
using dtype = typename RT::dtype;
|
||||
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.width; j++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.height; i++) {
|
||||
dst.tiles[i][j].data.thread_elements()[0] = op::template op<dtype>(a.tiles[i][j].data.thread_elements()[0], b.tiles[i][j].data.thread_elements()[0], col_values[j][0]);
|
||||
dst.tiles[i][j].data.thread_elements()[1] = op::template op<dtype>(a.tiles[i][j].data.thread_elements()[1], b.tiles[i][j].data.thread_elements()[1], col_values[j][0]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
|
||||
|
||||
// All of the annoying qualifiers *should* be automatically inferred during compile-time.
|
||||
// So, syntax should just be mittens::add_row(tile, colvec);
|
||||
|
||||
/**
|
||||
* @brief Sets all elements of a tile to zero.
|
||||
*
|
||||
* @tparam RT Tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
*/
|
||||
template<typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
zero(thread RT &dst) {
|
||||
unary_map<base_ops::zero, RT>(dst, dst);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of a tile to one.
|
||||
*
|
||||
* @tparam RT Tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
*/
|
||||
template<typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
one(thread RT &dst) {
|
||||
unary_map<base_ops::one, RT>(dst, dst);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of a tile to positive infinity.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
*/
|
||||
template<typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
pos_infty(thread RT &dst) {
|
||||
unary_map<base_ops::pos_infty, RT>(dst, dst);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of a tile to negative infinity.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
*/
|
||||
template<typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
neg_infty(thread RT &dst) {
|
||||
unary_map<base_ops::neg_infty, RT>(dst, dst);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Applies the exponential function to each element of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the exponential function on.
|
||||
*/
|
||||
template<typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
exp(thread RT &dst, thread const RT &src) {
|
||||
unary_map<base_ops::exp, RT>(dst, src);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the exponential function to each element of a tile, in base 2.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the exponential function on.
|
||||
*/
|
||||
template<typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
exp2(thread RT &dst, thread const RT &src) {
|
||||
unary_map<base_ops::exp2, RT>(dst, src);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the natural logarithm function to each element of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the natural logarithm function on.
|
||||
*/
|
||||
template<typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
log(thread RT &dst, thread const RT &src) {
|
||||
unary_map<base_ops::log, RT>(dst, src);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the absolute value function to each element of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the absolute value function on.
|
||||
*/
|
||||
template<typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
abs(thread RT &dst, thread const RT &src) {
|
||||
unary_map<base_ops::abs, RT>(dst, src);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the rectified linear unit (ReLU) function to each element of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the ReLU function on.
|
||||
*/
|
||||
template<typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
relu(thread RT &dst, thread const RT &src) {
|
||||
unary_map<base_ops::relu, RT>(dst, src);
|
||||
}
|
||||
/**
|
||||
* @brief Copies the elements from one tile to another.
|
||||
*
|
||||
* @tparam T Destination tile type.
|
||||
* @tparam U Source tile type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to copy from.
|
||||
*/
|
||||
template<typename RT, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
copy(thread RT &dst, thread const U &src) {
|
||||
bin_map<base_ops::copy2, RT>(dst, dst, src);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Applies the max operation element-wise between two tiles or a tile and a scalar.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam U Second operand type, which can be a tile or a scalar.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param lhs[in] Left-hand side source tile for the operation.
|
||||
* @param rhs[in] Right-hand side source tile or scalar for the operation.
|
||||
*/
|
||||
template<typename RT, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
max(thread RT &dst, thread const RT &lhs, thread const U &rhs) {
|
||||
bin_map<base_ops::max, RT>(dst, lhs, rhs);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the min operation element-wise between two tiles or a tile and a scalar.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam U Second operand type, which can be a tile or a scalar.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param lhs[in] Left-hand side source tile for the operation.
|
||||
* @param rhs[in] Right-hand side source tile or scalar for the operation.
|
||||
*/
|
||||
template<typename RT, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
min(thread RT &dst, thread const RT &lhs, thread const U &rhs) {
|
||||
bin_map<base_ops::min, RT>(dst, lhs, rhs);
|
||||
}
|
||||
/**
|
||||
* @brief Adds two tiles element-wise or adds a scalar to each element of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam U Second operand type, which can be a tile or a scalar.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param lhs[in] Left-hand side source tile for the addition.
|
||||
* @param rhs[in] Right-hand side source tile or scalar for the addition.
|
||||
*/
|
||||
template<typename RT, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
add(thread RT &dst, thread const RT &lhs, thread const U &rhs) {
|
||||
bin_map<base_ops::sum, RT>(dst, lhs, rhs);
|
||||
}
|
||||
/**
|
||||
* @brief Subtracts two tiles element-wise or subtracts a scalar from each element of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam U Second operand type, which can be a tile or a scalar.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param lhs[in] Left-hand side source tile for the subtraction.
|
||||
* @param rhs[in] Right-hand side source tile or scalar for the subtraction.
|
||||
*/
|
||||
template<typename RT, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
sub(thread RT &dst, const thread RT &lhs, thread const U &rhs) {
|
||||
bin_map<base_ops::sub, RT>(dst, lhs, rhs);
|
||||
}
|
||||
/**
|
||||
* @brief Multiplies two tiles element-wise or multiplies each element of a tile by a scalar.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam U Second operand type, which can be a tile or a scalar.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param lhs[in] Left-hand side source tile for the multiplication.
|
||||
* @param rhs[in] Right-hand side source tile or scalar for the multiplication.
|
||||
*/
|
||||
template<typename RT, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
mul(thread RT &dst, thread const RT &lhs, thread const U &rhs) {
|
||||
bin_map<base_ops::mul, RT>(dst, lhs, rhs);
|
||||
}
|
||||
/**
|
||||
* @brief Divides two tiles element-wise or divides each element of a tile by a scalar.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam U Second operand type, which can be a tile or a scalar.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param lhs[in] Left-hand side source tile for the division.
|
||||
* @param rhs[in] Right-hand side source tile or scalar for the division.
|
||||
*/
|
||||
template<typename RT, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>(), void>::type
|
||||
div(thread RT &dst, thread const RT &lhs, thread const U &rhs) {
|
||||
bin_map<base_ops::div, RT>(dst, lhs, rhs);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Adds row values to each row of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the addition on.
|
||||
* @param row_values[in] Column vector containing values to add to each row.
|
||||
*/
|
||||
template<typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
add_row(thread RT &dst, thread const RT &src, thread const RV &row_values) {
|
||||
row_map<base_ops::sum, RT, RV>(dst, src, row_values);
|
||||
}
|
||||
/**
|
||||
* @brief Subtracts row values from each row of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the subtraction on.
|
||||
* @param row_values[in] Column vector containing values to subtract from each row.
|
||||
*/
|
||||
template<typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
sub_row(thread RT &dst, thread const RT &src, thread const RV &row_values) {
|
||||
row_map<base_ops::sub, RT, RV>(dst, src, row_values);
|
||||
// using T4 = typename base_types::packing<typename RT::dtype>::packed_four;
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < RT::height; i++) {
|
||||
// // #pragma clang loop unroll(full)
|
||||
// // for(int j = 0; j < RT::width; j+=2) {
|
||||
// // T4 val = op::template op<T4>({src.tiles[i][j].data.thread_elements()[0],
|
||||
// // src.tiles[i][j].data.thread_elements()[1],
|
||||
// // src.tiles[i][j+1].data.thread_elements()[0],
|
||||
// // src.tiles[i][j+1].data.thread_elements()[1],},
|
||||
// // {row_values[i][0], row_values[i][0],row_values[i][0], row_values[i][0]});
|
||||
// //
|
||||
// // dst.tiles[i][j].data.thread_elements()[0] = val[0];
|
||||
// // dst.tiles[i][j].data.thread_elements()[1] = val[1];
|
||||
// // dst.tiles[i][j+1].data.thread_elements()[0] = val[2];
|
||||
// // dst.tiles[i][j+1].data.thread_elements()[1] = val[3];
|
||||
// // }
|
||||
//
|
||||
// // #pragma clang loop unroll(full)
|
||||
// // for(int j = 0; j < RT::width; j++) {
|
||||
// // T2 val = op::template op<T2>({src.tiles[i][j].data.thread_elements()[0],
|
||||
// // src.tiles[i][j].data.thread_elements()[1]},
|
||||
// // {row_values[i][0], row_values[i][0]});
|
||||
// //
|
||||
// // dst.tiles[i][j].data.thread_elements()[0] = val[0];
|
||||
// // dst.tiles[i][j].data.thread_elements()[1] = val[1];
|
||||
// // }
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < RT::width; j+=2) {
|
||||
// T4 val = T4(src.tiles[i][j].data.thread_elements()[0],
|
||||
// src.tiles[i][j].data.thread_elements()[1],
|
||||
// src.tiles[i][j+1].data.thread_elements()[0],
|
||||
// src.tiles[i][j+1].data.thread_elements()[1]) - T4(row_values[i][0], row_values[i][0], row_values[i][0], row_values[i][0]);
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = val[0];
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = val[1];
|
||||
// dst.tiles[i][j+1].data.thread_elements()[0] = val[2];
|
||||
// dst.tiles[i][j+1].data.thread_elements()[1] = val[3];
|
||||
// }
|
||||
// }
|
||||
}
|
||||
/**
|
||||
* @brief Multiplies each row of a tile by row values.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the multiplication on.
|
||||
* @param row_values[in] Column vector containing values to multiply each row by.
|
||||
*/
|
||||
template<typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
mul_row(thread RT &dst, thread const RT &src, thread const RV &row_values) {
|
||||
// using T = typename RT::T;
|
||||
// using T2 = typename RT::T2;
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < RT::height; i++) {
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 0; j < RT::width; j++) {
|
||||
//// T s1 = src.tiles[i][j].data.thread_elements()[0];
|
||||
//// T v1 = row_values[i][0];
|
||||
//// dst.tiles[i][j].data.thread_elements()[0] = s1 * v1;
|
||||
//// T s2 = src.tiles[i][j].data.thread_elements()[1];
|
||||
//// T v2 = row_values[i][1];
|
||||
//// dst.tiles[i][j].data.thread_elements()[1] = s2 * v2;
|
||||
//
|
||||
//
|
||||
//// dst.tiles[i][j].data.thread_elements()[0] = op::template op<T>(src.tiles[i][j].data.thread_elements()[0], row_values[i][0]);
|
||||
//// dst.tiles[i][j].data.thread_elements()[1] = op::template op<T>(src.tiles[i][j].data.thread_elements()[1], row_values[i][0]);
|
||||
// T2 val = op::template op<T2>({src.tiles[i][j].data.thread_elements()[0], row_values[i][0]);
|
||||
// dst.tiles[i][j].data.thread_elements()[0] = op::template op<T>(src.tiles[i][j].data.thread_elements()[0], row_values[i][0]);
|
||||
// dst.tiles[i][j].data.thread_elements()[1] = op::template op<T>(src.tiles[i][j].data.thread_elements()[1], row_values[i][0]);
|
||||
// }
|
||||
// }
|
||||
row_map<base_ops::mul, RT, RV>(dst, src, row_values);
|
||||
}
|
||||
/**
|
||||
* @brief Divides each row of a tile by row values.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the division on.
|
||||
* @param row_values[in] Column vector containing values to divide each row by.
|
||||
*/
|
||||
template<typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
div_row(thread RT &dst, thread const RT &src, thread const RV &row_values) {
|
||||
row_map<base_ops::div, RT, RV>(dst, src, row_values);
|
||||
}
|
||||
/**
|
||||
* @brief Broadcast a vector into into a tile's rows.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param row_values[in] Column vector containing values to broadcast into rows.
|
||||
*/
|
||||
template<typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
broadcast_row(thread RT &dst, thread const RV &row_values) {
|
||||
row_map<base_ops::copy2, RT, RV>(dst, dst, row_values);
|
||||
}
|
||||
|
||||
|
||||
// col maps
|
||||
/**
|
||||
* @brief Adds column values to each column of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the addition on.
|
||||
* @param col_values[in] Row vector containing values to add to each column.
|
||||
*/
|
||||
template<typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
add_col(thread RT &dst, thread const RT &src, thread const RV &col_values) {
|
||||
col_map<base_ops::sum, RT, RV>(dst, src, col_values);
|
||||
}
|
||||
/**
|
||||
* @brief Subtracts column values from each column of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the subtraction on.
|
||||
* @param col_values[in] Row vector containing values to subtract from each column.
|
||||
*/
|
||||
template<typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
sub_col(thread RT &dst, thread const RT &src, thread const RV &col_values) {
|
||||
col_map<base_ops::sub, RT, RV>(dst, src, col_values);
|
||||
}
|
||||
/**
|
||||
* @brief Multiplies each column of a tile by column values.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the multiplication on.
|
||||
* @param col_values[in] Row vector containing values to multiply each column by.
|
||||
*/
|
||||
template<typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
mul_col(thread RT &dst, thread const RT &src, thread const RV &col_values) {
|
||||
col_map<base_ops::mul, RT, RV>(dst, src, col_values);
|
||||
}
|
||||
/**
|
||||
* @brief Divides each column of a tile by column values.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the division on.
|
||||
* @param col_values[in] Row vector containing values to divide each column by.
|
||||
*/
|
||||
template<typename RT, typename RV>
|
||||
static METAL_FUNC void div_col(thread RT &dst, thread const RT &src, thread const RV &col_values) {
|
||||
col_map<base_ops::div, RT, RV>(dst, src, col_values);
|
||||
}
|
||||
/**
|
||||
* @brief Broadcast a vector into into a tile's columns.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param row_values[in] Row vector containing values to broadcast into cols.
|
||||
*/
|
||||
template<typename RT, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
broadcast_col(thread RT &dst, thread const RV &col_values) {
|
||||
col_map<base_ops::copy2, RT, RV>(dst, dst, col_values);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,214 @@
|
||||
#pragma once // doneington
|
||||
|
||||
#include <metal_stdlib>
|
||||
#include "../../../../types/types.metal"
|
||||
#include "../../../../common/common.metal"
|
||||
namespace mittens {
|
||||
|
||||
template <typename R, typename T, typename U, typename V,
|
||||
typename l1, typename l2, typename l3, typename l4>
|
||||
METAL_FUNC static void mma_base(thread rt_base<R, l1>& d,
|
||||
thread rt_base<T, l2>& a,
|
||||
thread rt_base<U, l3>& b,
|
||||
thread rt_base<V, l4>& c) {
|
||||
metal::simdgroup_multiply_accumulate(d.data, a.data, b.data, c.data);
|
||||
}
|
||||
|
||||
template <typename R, typename T, typename U,
|
||||
typename l1, typename l2, typename l3>
|
||||
METAL_FUNC static void mm_base(thread rt_base<R, l1>& d,
|
||||
thread rt_base<T, l2>& a,
|
||||
thread rt_base<U, l3>& b) {
|
||||
metal::simdgroup_multiply(d.data, a.data, b.data);
|
||||
}
|
||||
|
||||
namespace meta {
|
||||
template<typename R, typename T, typename U, int N, int K, int M>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::base_types::isT1Type<R>() && ducks::base_types::isT1Type<T>() && ducks::base_types::isT1Type<U>(), void>::type
|
||||
mma_AB_unroll_inner(int k, int n, int m,
|
||||
thread rt<R, N, M, ducks::rt_layout::row>* d,
|
||||
thread rt<T, N, K, ducks::rt_layout::row>* a,
|
||||
thread rt<U, K, M, ducks::rt_layout::row>* b) {
|
||||
mma_base(
|
||||
d->tiles[n][m],
|
||||
a->tiles[n][k],
|
||||
b->tiles[k][m],
|
||||
d->tiles[n][m]
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
template<typename R, typename T, typename U, typename V, int N, int K, int M>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::base_types::isT1Type<R>() && ducks::base_types::isT1Type<T>() && ducks::base_types::isT1Type<U>() && ducks::base_types::isT1Type<V>(), void>::type
|
||||
mma_AB_unroll(int n, int m,
|
||||
thread rt<R, N, M, ducks::rt_layout::row>* d,
|
||||
thread rt<T, N, K, ducks::rt_layout::row>* a,
|
||||
thread rt<U, K, M, ducks::rt_layout::row>* b,
|
||||
thread rt<V, N, M, ducks::rt_layout::row>* c) {
|
||||
mma_base(
|
||||
d->tiles[n][m],
|
||||
a->tiles[n][0],
|
||||
b->tiles[0][m],
|
||||
c->tiles[n][m]
|
||||
);
|
||||
meta::unroll_i_in_range<1, K/TILE_DIM, 1>::run(meta::mma_AB_unroll_inner<R, T, U, N, K, M>, n, m, d, a, b);
|
||||
}
|
||||
|
||||
template<typename R, typename T, typename U, int N, int K, int M>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::base_types::isT1Type<R>() && ducks::base_types::isT1Type<T>() && ducks::base_types::isT1Type<U>(), void>::type
|
||||
mm_AB_unroll(int n, int m,
|
||||
thread rt<R, N, M, ducks::rt_layout::row>* d,
|
||||
thread rt<T, N, K, ducks::rt_layout::row>* a,
|
||||
thread rt<U, K, M, ducks::rt_layout::row>* b) {
|
||||
mm_base(
|
||||
d->tiles[n][m],
|
||||
a->tiles[n][0],
|
||||
b->tiles[0][m]
|
||||
);
|
||||
meta::unroll_i_in_range<1, K/TILE_DIM, 1>::run(meta::mma_AB_unroll_inner<R, T, U, N, K, M>, n, m, d, a, b);
|
||||
}
|
||||
}
|
||||
|
||||
template<typename R, typename T, typename U, typename V, int N, int K, int M>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::base_types::isT1Type<R>() && ducks::base_types::isT1Type<T>() && ducks::base_types::isT1Type<U>() && ducks::base_types::isT1Type<V>(), void>::type
|
||||
mma_AB(thread rt<R, N, M, ducks::rt_layout::row>& d,
|
||||
thread rt<T, N, K, ducks::rt_layout::row>& a,
|
||||
thread rt<U, K, M, ducks::rt_layout::row>& b,
|
||||
thread rt<V, N, M, ducks::rt_layout::row>& c) {
|
||||
meta::unroll_i_j_in_range<0, N/TILE_DIM, 1, 0, M/TILE_DIM, 1>::run(meta::mma_AB_unroll<R, T, U, V, N, K, M>, &d, &a, &b, &c);
|
||||
}
|
||||
|
||||
template<typename R, typename T, typename U, int N, int K, int M>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::base_types::isT1Type<R>() && ducks::base_types::isT1Type<T>() && ducks::base_types::isT1Type<U>(), void>::type
|
||||
mm_AB(thread rt<R, N, M, ducks::rt_layout::row>& d,
|
||||
thread rt<T, N, K, ducks::rt_layout::row>& a,
|
||||
thread rt<U, K, M, ducks::rt_layout::row>& b) {
|
||||
meta::unroll_i_j_in_range<0, N/TILE_DIM, 1, 0, M/TILE_DIM, 1>::run(meta::mm_AB_unroll<R, T, U, N, K, M>, &d, &a, &b);
|
||||
}
|
||||
|
||||
namespace meta {
|
||||
template<typename R, typename T, typename U, int N, int K, int M>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::base_types::isT1Type<R>() && ducks::base_types::isT1Type<T>() && ducks::base_types::isT1Type<U>(), void>::type
|
||||
mma_ABt_unroll_inner(int k, int n, int m,
|
||||
thread rt<R, N, M, ducks::rt_layout::row>* d,
|
||||
thread rt<T, N, K, ducks::rt_layout::row>* a,
|
||||
thread rt<U, M, K, ducks::rt_layout::col>* b) {
|
||||
mma_base(
|
||||
d->tiles[n][m],
|
||||
a->tiles[n][k],
|
||||
b->tiles[m][k],
|
||||
d->tiles[n][m]
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
template<typename R, typename T, typename U, typename V, int N, int K, int M>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::base_types::isT1Type<R>() && ducks::base_types::isT1Type<T>() && ducks::base_types::isT1Type<U>() && ducks::base_types::isT1Type<V>(), void>::type
|
||||
mma_ABt_unroll(int n, int m,
|
||||
thread rt<R, N, M, ducks::rt_layout::row>* d,
|
||||
thread rt<T, N, K, ducks::rt_layout::row>* a,
|
||||
thread rt<U, M, K, ducks::rt_layout::col>* b,
|
||||
thread rt<V, N, M, ducks::rt_layout::row>* c) {
|
||||
mma_base(
|
||||
d->tiles[n][m],
|
||||
a->tiles[n][0],
|
||||
b->tiles[m][0],
|
||||
c->tiles[n][m]
|
||||
);
|
||||
meta::unroll_i_in_range<1, K/TILE_DIM, 1>::run(meta::mma_ABt_unroll_inner<R, T, U, N, K, M>, n, m, d, a, b);
|
||||
}
|
||||
|
||||
template<typename R, typename T, typename U, int N, int K, int M>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::base_types::isT1Type<R>() && ducks::base_types::isT1Type<T>() && ducks::base_types::isT1Type<U>(), void>::type
|
||||
mm_ABt_unroll(int n, int m,
|
||||
thread rt<R, N, M, ducks::rt_layout::row>* d,
|
||||
thread rt<T, N, K, ducks::rt_layout::row>* a,
|
||||
thread rt<U, M, K, ducks::rt_layout::col>* b) {
|
||||
mm_base(
|
||||
d->tiles[n][m],
|
||||
a->tiles[n][0],
|
||||
b->tiles[m][0]
|
||||
);
|
||||
meta::unroll_i_in_range<1, K/TILE_DIM, 1>::run(meta::mma_ABt_unroll_inner<R, T, U, N, K, M>, n, m, d, a, b);
|
||||
}
|
||||
}
|
||||
|
||||
template<typename R, typename T, typename U, typename V, int N, int K, int M>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::base_types::isT1Type<R>() && ducks::base_types::isT1Type<T>() && ducks::base_types::isT1Type<U>() && ducks::base_types::isT1Type<V>(), void>::type
|
||||
mma_ABt(thread rt<R, N, M, ducks::rt_layout::row>& d,
|
||||
thread rt<T, N, K, ducks::rt_layout::row>& a,
|
||||
thread rt<U, M, K, ducks::rt_layout::col>& b,
|
||||
thread rt<V, N, M, ducks::rt_layout::row>& c) {
|
||||
meta::unroll_i_j_in_range<0, N/TILE_DIM, 1, 0, M/TILE_DIM, 1>::run(meta::mma_ABt_unroll<R, T, U, V, N, K, M>, &d, &a, &b, &c);
|
||||
}
|
||||
|
||||
template<typename R, typename T, typename U, int N, int K, int M>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::base_types::isT1Type<R>() && ducks::base_types::isT1Type<T>() && ducks::base_types::isT1Type<U>(), void>::type
|
||||
mm_ABt(thread rt<R, N, M, ducks::rt_layout::row>& d,
|
||||
thread rt<T, N, K, ducks::rt_layout::row>& a,
|
||||
thread rt<U, M, K, ducks::rt_layout::col>& b) {
|
||||
meta::unroll_i_j_in_range<0, N/TILE_DIM, 1, 0, M/TILE_DIM, 1>::run(meta::mm_ABt_unroll<R, T, U, N, K, M>, &d, &a, &b);
|
||||
}
|
||||
|
||||
template<typename R, typename T, typename U, typename V, int N, int K, int M>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::base_types::isT1Type<R>() && ducks::base_types::isT1Type<T>() && ducks::base_types::isT1Type<U>() && ducks::base_types::isT1Type<V>(), void>::type
|
||||
mma_AtB(thread rt<R, N, M, ducks::rt_layout::row>& d,
|
||||
thread rt<T, K, N, ducks::rt_layout::col>& a,
|
||||
thread rt<U, K, M, ducks::rt_layout::row>& b,
|
||||
thread rt<V, N, M, ducks::rt_layout::row>& c) {
|
||||
#pragma clang loop unroll(full)
|
||||
for (int n = 0; n < N / TILE_DIM; n++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for (int m = 0; m < M / TILE_DIM; m++) {
|
||||
mma_base(
|
||||
d.tiles[n][m],
|
||||
a.tiles[0][n],
|
||||
b.tiles[0][m],
|
||||
c.tiles[n][m]
|
||||
);
|
||||
#pragma clang loop unroll(full)
|
||||
for (int k = 1; k < K / TILE_DIM; k++) {
|
||||
mma_base(
|
||||
d.tiles[n][m],
|
||||
a.tiles[k][n],
|
||||
b.tiles[k][m],
|
||||
d.tiles[n][m]
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template<typename R, typename T, typename U, typename V, int N, int K, int M>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::base_types::isT1Type<R>() && ducks::base_types::isT1Type<T>() && ducks::base_types::isT1Type<U>() && ducks::base_types::isT1Type<V>(), void>::type
|
||||
mma_AtBt(thread rt<R, N, M, ducks::rt_layout::row>& d,
|
||||
thread rt<T, K, N, ducks::rt_layout::col>& a,
|
||||
thread rt<U, M, K, ducks::rt_layout::col>& b,
|
||||
thread rt<V, N, M, ducks::rt_layout::row>& c) {
|
||||
#pragma clang loop unroll(full)
|
||||
for (int n = 0; n < N / TILE_DIM; n++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for (int m = 0; m < M / TILE_DIM; m++) {
|
||||
mma_base(
|
||||
d.tiles[n][m],
|
||||
a.tiles[0][n],
|
||||
b.tiles[m][0],
|
||||
c.tiles[n][m]
|
||||
);
|
||||
#pragma clang loop unroll(full)
|
||||
for (int k = 1; k < K / TILE_DIM; k++) {
|
||||
mma_base(
|
||||
d.tiles[n][m],
|
||||
a.tiles[k][n],
|
||||
b.tiles[m][k],
|
||||
d.tiles[n][m]
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,636 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Reduction operations mapping tiles to vectors.
|
||||
*/
|
||||
|
||||
#pragma once //doneington (but register col layotus)
|
||||
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
|
||||
namespace meta {
|
||||
|
||||
//template<typename op, typename RT>
|
||||
//static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>(), void>::type
|
||||
//row_reduce_unroll_inner(int i, thread const RT *src, thread typename RT::T& accum_thread) {
|
||||
// accum_thread = op::template op<typename RT::T>(accum_thread, src->tiles[i][0].data.thread_elements()[0]);
|
||||
// accum_thread = op::template op<typename RT::T>(accum_thread, src->tiles[i][0].data.thread_elements()[1]);
|
||||
//}
|
||||
//
|
||||
//template<typename op, typename RV, typename RT, bool reset>
|
||||
//static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
//row_reduce_unroll(int i, thread RV *row_accum, thread const RT *src, thread const RV *src_accum, const short leader) {
|
||||
// using T = typename RV::T;
|
||||
// T accum_thread = op::template op<T>(src->tiles[i][0].data.thread_elements()[0], src->tiles[i][0].data.thread_elements()[1]);
|
||||
//
|
||||
// meta::unroll_i_in_range<1, RT::width, 1>::run(meta::row_reduce_unroll_inner<op, RT>, src, accum_thread);
|
||||
// accum_thread = op::template op<T>(accum_thread, shfl_down_sync<T>(accum_thread, 1));
|
||||
// accum_thread = op::template op<T>(accum_thread, shfl_down_sync<T>(accum_thread, 8));
|
||||
//
|
||||
// accum_thread = shfl_sync<T>(accum_thread, leader);
|
||||
//
|
||||
// if(reset) { (*row_accum)[i][0] = accum_thread; }
|
||||
// else { (*row_accum)[i][0] = op::template op<T>((*src_accum)[i][0], accum_thread); }
|
||||
//}
|
||||
|
||||
//template<typename op, typename RT>
|
||||
//static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>(), void>::type
|
||||
//row_reduce_unroll_inner(int i, thread const RT *src, thread typename RT::T2& accum_thread) {
|
||||
// accum_thread = op::template op<typename RT::T2>(accum_thread, {src->tiles[i][0].data.thread_elements()[0], src->tiles[i][0].data.thread_elements()[1]});
|
||||
//}
|
||||
|
||||
/*
|
||||
pragma clang loop unroll(full)
|
||||
for(int i = 0; i < src.height; i++) {
|
||||
T accum_thread = op::template op<T>(src.tiles[i][0].data.thread_elements()[0], src.tiles[i][0].data.thread_elements()[1]);
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 1; j < src.width; j++) {
|
||||
accum_thread = op::template op<T>(accum_thread, src.tiles[i][j].data.thread_elements()[0]);
|
||||
accum_thread = op::template op<T>(accum_thread, src.tiles[i][j].data.thread_elements()[1]);
|
||||
}
|
||||
accum_thread = op::template op<T>(accum_thread, shfl_down_sync<T>(accum_thread, 1));
|
||||
accum_thread = op::template op<T>(accum_thread, shfl_down_sync<T>(accum_thread, 8));
|
||||
|
||||
accum_thread = shfl_sync<T>(accum_thread, leader);
|
||||
|
||||
if(reset) { row_accum[i][0] = accum_thread; }
|
||||
else { row_accum[i][0] = op::template op<T>(src_accum[i][0], accum_thread); }
|
||||
}
|
||||
*/
|
||||
|
||||
template<typename op, typename RV, typename RT, bool reset>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_reduce_unroll(int i, thread RV *row_accum, thread const RT *src, thread const RV *src_accum, const short leader) {
|
||||
using T = typename RV::T;
|
||||
using T2 = typename RV::T2;
|
||||
T accum_thread = op::template op<T>(src->tiles[i][0].data.thread_elements()[0], src->tiles[i][0].data.thread_elements()[1]);
|
||||
for(int j = 1; j < src->width; j++) {
|
||||
accum_thread = op::template op<T>(accum_thread, src->tiles[i][j].data.thread_elements()[0]);
|
||||
accum_thread = op::template op<T>(accum_thread, src->tiles[i][j].data.thread_elements()[1]);
|
||||
}
|
||||
|
||||
T shfl_val = shfl_down_sync<T>(accum_thread, 1);
|
||||
accum_thread = op::template op<T>(accum_thread, shfl_val);
|
||||
shfl_val = shfl_down_sync<T>(accum_thread, 8);
|
||||
accum_thread = op::template op<T>(accum_thread, shfl_val);
|
||||
|
||||
accum_thread = shfl_sync<T>(accum_thread, leader);
|
||||
|
||||
if(reset) {
|
||||
(*row_accum)[i][0] = accum_thread;
|
||||
}
|
||||
else {
|
||||
(*row_accum)[i][0] = op::template op<T>((*src_accum)[i][0], accum_thread);;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
/**
|
||||
* @brief Perform a row-wise reduction on a matrix in row-major layout.
|
||||
*
|
||||
* This function template performs a parallel reduction across the rows of a matrix using a specified operation.
|
||||
* It leverages warp shuffle functions for efficient intra-warp communication.
|
||||
*
|
||||
* @tparam op The operation to be applied for reduction.
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type with row layout.
|
||||
* @tparam reset A boolean flag indicating whether to reset the accumulator (ignore src_accum) or not.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when reset is false.
|
||||
*/
|
||||
template<typename op, typename RV, typename RT, bool reset>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_reduce(thread RV &row_accum, thread const RT &src, thread const RV &src_accum, const short laneid) {
|
||||
static_assert(ducks::is_ortho_layout<typename RV::layout>(), "rv must be ortho for row RT");
|
||||
static_assert(metal::is_same_v<typename RV::dtype, typename RT::dtype>, "rv and rt must be the same type"); // compatible type
|
||||
static_assert(RV::outer_dim == RT::height, "rv and rt dims don't match"); // compatible size
|
||||
using T = typename RV::T;
|
||||
using T2 = typename RV::T2;
|
||||
const short leader = (laneid / 16) * 16 + ((laneid / 2) % 4) * 2;
|
||||
|
||||
// constexpr const uint32_t COL_0 = 0x00550055;
|
||||
// constexpr const uint32_t COL_1 = 0x00AA00AA;
|
||||
// constexpr const uint32_t COL_2 = 0x55005500;
|
||||
// constexpr const uint32_t COL_3 = 0xAA00AA00;
|
||||
//
|
||||
// constexpr const uint32_t COL_0_2 = COL_0 | COL_2;
|
||||
// constexpr const uint32_t COL_0_1 = COL_0 | COL_1;
|
||||
// constexpr const uint32_t COL_2_3 = COL_2 | COL_3;
|
||||
// const ushort src_lane1 = laneid + ((COL_0_2 >> laneid) & 1) * 1 + ((COL_1 >> laneid) & 1) * 7 - ((COL_3 >> laneid) & 1) * 9;
|
||||
// const ushort src_lane2 = laneid + ((COL_0_1 >> laneid) & 1) * 8 - ((COL_2_3 >> laneid) & 1) * 8;
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < src.height; i++) {
|
||||
// T accum_thread = op::template op<T>(src.tiles[i][0].data.thread_elements()[0], src.tiles[i][0].data.thread_elements()[1]);
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 1; j < src.width; j++) {
|
||||
// accum_thread = op::template op<T>(accum_thread, src.tiles[i][j].data.thread_elements()[0]);
|
||||
// accum_thread = op::template op<T>(accum_thread, src.tiles[i][j].data.thread_elements()[1]);
|
||||
// }
|
||||
// accum_thread = op::template op<T>(accum_thread, shfl_sync<T>(accum_thread, src_lane1));
|
||||
// accum_thread = op::template op<T>(accum_thread, shfl_sync<T>(accum_thread, src_lane2));
|
||||
//
|
||||
//
|
||||
// if(reset) { row_accum[i][0] = accum_thread; }
|
||||
// else { row_accum[i][0] = op::template op<T>(src_accum[i][0], accum_thread); }
|
||||
// }
|
||||
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < src.height; i++) {
|
||||
// T accum_thread = op::template op<T>(src.tiles[i][0].data.thread_elements()[0], src.tiles[i][0].data.thread_elements()[1]);
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 1; j < src.width; j++) {
|
||||
// accum_thread = op::template op<T>(accum_thread, src.tiles[i][j].data.thread_elements()[0]);
|
||||
// accum_thread = op::template op<T>(accum_thread, src.tiles[i][j].data.thread_elements()[1]);
|
||||
// }
|
||||
// accum_thread = op::template op<T>(accum_thread, shfl_down_sync<T>(accum_thread, 1));
|
||||
// accum_thread = op::template op<T>(accum_thread, shfl_down_sync<T>(accum_thread, 8));
|
||||
//
|
||||
// accum_thread = shfl_sync<T>(accum_thread, leader);
|
||||
//
|
||||
// if(reset) { row_accum[i][0] = accum_thread; }
|
||||
// else { row_accum[i][0] = op::template op<T>(src_accum[i][0], accum_thread); }
|
||||
// }
|
||||
|
||||
meta::unroll_i_in_range<0, RT::height, 1>::run(meta::row_reduce_unroll<op, RV, RT, reset>, &row_accum, &src, &src_accum, leader);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Perform a row-wise reduction on a matrix in row-major layout.
|
||||
*
|
||||
* This function template performs a parallel reduction across the rows of a matrix using a specified operation.
|
||||
* It leverages warp shuffle functions for efficient intra-warp communication.
|
||||
*
|
||||
* @tparam op The operation to be applied for reduction.
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type with row layout.
|
||||
* @tparam reset A boolean flag indicating whether to reset the accumulator (ignore src_accum) or not.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when reset is false.
|
||||
*/
|
||||
template<typename op, typename RV, typename RT, bool reset>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_reduce(thread RV &row_accum, thread const RT &src, thread const RV &src_accum, const short laneid) {
|
||||
static_assert(ducks::is_align_layout<typename RV::layout>(), "rv must be align for row RT");
|
||||
static_assert(metal::is_same_v<typename RV::dtype, typename RT::dtype>, "rv and rt must be the same type"); // compatible type
|
||||
static_assert(RV::outer_dim == RT::height, "rv and rt dims don't match"); // compatible size
|
||||
|
||||
using T = typename RV::T;
|
||||
using T2 = typename RV::T2;
|
||||
|
||||
const int leader = (laneid % 2) + ((laneid / 8) % 2) * 8;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < src.height; i++) {
|
||||
T2 accum_thread = {src.tiles[i][0].data.thread_elements()[0], src.tiles[i][0].data.thread_elements()[1]};
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 1; j < src.width; j++) {
|
||||
accum_thread = op::template op<T2>(accum_thread, {src.tiles[i][j].data.thread_elements()[0], src.tiles[i][j].data.thread_elements()[1]});
|
||||
}
|
||||
// Now we need to do a lil shuffle to make everyone happy.
|
||||
|
||||
accum_thread = op::template op<T2>(accum_thread, shfl_down_sync<T2>(accum_thread, 2));
|
||||
accum_thread = op::template op<T2>(accum_thread, shfl_down_sync<T2>(accum_thread, 4));
|
||||
accum_thread = op::template op<T2>(accum_thread, shfl_down_sync<T2>(accum_thread, 16));
|
||||
|
||||
accum_thread = shfl_sync<T2>(accum_thread, leader);
|
||||
|
||||
if(reset) {
|
||||
row_accum[i][0] = accum_thread[0];
|
||||
row_accum[i][1] = accum_thread[1];
|
||||
}
|
||||
else {
|
||||
row_accum[i][0] = op::template op<T>(row_accum[i][0], accum_thread[0]);
|
||||
row_accum[i][1] = op::template op<T>(row_accum[i][1], accum_thread[1]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* @brief Perform a column-wise reduction on a matrix in row-major layout.
|
||||
*
|
||||
* This function template performs a parallel reduction across the columns of a matrix using a specified operation.
|
||||
* It leverages warp shuffle functions for efficient intra-warp communication and is optimized for row-major matrices.
|
||||
*
|
||||
* @tparam op The operation to be applied for reduction.
|
||||
* @tparam V The vector type for the column accumulator.
|
||||
* @tparam T The matrix type with row layout.
|
||||
* @tparam reset A boolean flag indicating whether to reset the accumulator (ignore src_accum) or not.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when reset is false.
|
||||
*/
|
||||
template<typename op, typename RV, typename RT, bool reset>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_reduce(thread RV &col_accum, thread const RT &src, thread const RV &src_accum, const ushort laneid) {
|
||||
static_assert(ducks::is_align_layout<typename RV::layout>(), "rv must be align layout");
|
||||
static_assert(metal::is_same_v<typename RV::dtype, typename RT::dtype>, "rt and rv must be same type"); // compatible type
|
||||
static_assert(RV::outer_dim == RT::width, "rv and rt dims don't match"); // compatible size
|
||||
|
||||
using dtype = typename RV::dtype;
|
||||
using T2 = typename base_types::packing<dtype>::packed_type;
|
||||
|
||||
const int leader = (laneid % 2) + ((laneid / 8) % 2) * 8;
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < src.width; j++) {
|
||||
// dtype accum_left_cols = src.tiles[0][j].data.thread_elements()[0];
|
||||
// dtype accum_right_cols = src.tiles[0][j].data.thread_elements()[1];
|
||||
T2 accum_cols = {src.tiles[0][j].data.thread_elements()[0], src.tiles[0][j].data.thread_elements()[1]};
|
||||
// dtype accum_right_cols = src.tiles[0][j].data.thread_elements()[1];
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 1; i < src.height; i++) {
|
||||
// accum_left_cols = op::template op<dtype>(accum_left_cols , src.tiles[i][j].data.thread_elements()[0]);
|
||||
// accum_right_cols = op::template op<dtype>(accum_right_cols, src.tiles[i][j].data.thread_elements()[1]);
|
||||
accum_cols = op::template op<T2>(accum_cols, {src.tiles[i][j].data.thread_elements()[0], src.tiles[i][j].data.thread_elements()[1]});
|
||||
}
|
||||
|
||||
// accum_left_cols = op::template op<dtype>(accum_left_cols, shfl_down_sync<dtype>(accum_left_cols, 2));
|
||||
// accum_left_cols = op::template op<dtype>(accum_left_cols, shfl_down_sync<dtype>(accum_left_cols, 4));
|
||||
// accum_left_cols = op::template op<dtype>(accum_left_cols, shfl_down_sync<dtype>(accum_left_cols, 16));
|
||||
|
||||
// accum_right_cols = op::template op<dtype>(accum_right_cols, shfl_down_sync<dtype>(accum_right_cols, 2));
|
||||
// accum_right_cols = op::template op<dtype>(accum_right_cols, shfl_down_sync<dtype>(accum_right_cols, 4));
|
||||
// accum_right_cols = op::template op<dtype>(accum_right_cols, shfl_down_sync<dtype>(accum_right_cols, 16));
|
||||
accum_cols = op::template op<T2>(accum_cols, shfl_down_sync<T2>(accum_cols, 2));
|
||||
accum_cols = op::template op<T2>(accum_cols, shfl_down_sync<T2>(accum_cols, 4));
|
||||
accum_cols = op::template op<T2>(accum_cols, shfl_down_sync<T2>(accum_cols, 16));
|
||||
|
||||
// accum_left_cols = shfl_sync<dtype>(accum_left_cols, leader);
|
||||
// accum_right_cols = shfl_sync<dtype>(accum_right_cols, leader);
|
||||
accum_cols = shfl_sync<T2>(accum_cols, leader);
|
||||
|
||||
|
||||
if(reset) {
|
||||
// col_accum[j][0] = accum_left_cols;
|
||||
// col_accum[j][1] = accum_right_cols;
|
||||
col_accum[j][0] = accum_cols[0];
|
||||
col_accum[j][1] = accum_cols[1];
|
||||
}
|
||||
else {
|
||||
// col_accum[j][0] = op::template op<dtype>(src_accum[j][0], accum_left_cols);
|
||||
// col_accum[j][1] = op::template op<dtype>(src_accum[j][1], accum_right_cols);
|
||||
col_accum[j][0] = op::template op<dtype>(src_accum[j][0], accum_cols[0]);
|
||||
col_accum[j][1] = op::template op<dtype>(src_accum[j][1], accum_cols[1]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Perform a column-wise reduction on a matrix in row-major layout.
|
||||
*
|
||||
* This function template performs a parallel reduction across the columns of a matrix using a specified operation.
|
||||
* It leverages warp shuffle functions for efficient intra-warp communication and is optimized for row-major matrices.
|
||||
*
|
||||
* @tparam op The operation to be applied for reduction.
|
||||
* @tparam V The vector type for the column accumulator.
|
||||
* @tparam T The matrix type with row layout.
|
||||
* @tparam reset A boolean flag indicating whether to reset the accumulator (ignore src_accum) or not.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when reset is false.
|
||||
*/
|
||||
template<typename op, typename RV, typename RT, bool reset>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_reduce(thread RV &col_accum, thread const RT &src, thread const RV &src_accum, const ushort laneid) {
|
||||
static_assert(ducks::is_ortho_layout<typename RV::layout>(), "rv must be ortho layout");
|
||||
static_assert(metal::is_same_v<typename RV::dtype, typename RT::dtype>, "rt and rv must be same type"); // compatible type
|
||||
static_assert(RV::outer_dim == RT::width, "rv and rt dims don't match"); // compatible size
|
||||
|
||||
using T = typename RV::T;
|
||||
using T2 = typename base_types::packing<T>::packed_type;
|
||||
|
||||
const int leader = (laneid / 16) * 16 + ((laneid / 2) % 4) * 2; // lololol
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < src.width; i++) {
|
||||
T accum_thread = op::template op<T>(src.tiles[0][i].data.thread_elements()[0], src.tiles[0][i].data.thread_elements()[1]);
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 1; j < src.height; j++) {
|
||||
accum_thread = op::template op<T>(accum_thread, src.tiles[j][i].data.thread_elements()[0]);
|
||||
accum_thread = op::template op<T>(accum_thread, src.tiles[j][i].data.thread_elements()[1]);
|
||||
}
|
||||
// Now we need to do a lil shuffle to make everyone happy.
|
||||
|
||||
accum_thread = op::template op<T>(accum_thread, shfl_down_sync<T>(accum_thread, 1));
|
||||
accum_thread = op::template op<T>(accum_thread, shfl_down_sync<T>(accum_thread, 8));
|
||||
|
||||
accum_thread = shfl_sync<T>(accum_thread, leader);
|
||||
|
||||
if(reset) {
|
||||
col_accum[i][0] = accum_thread;
|
||||
}
|
||||
else {
|
||||
col_accum[i][0] = op::template op<T>(col_accum[i][0], accum_thread);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
|
||||
// two-operand row reductions. (Accumulate and REPLACE.)
|
||||
/**
|
||||
* @brief Store the maximum of each row of the src register tile in the row_accum column vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_max(thread RV &row_accum, thread const RT &src, const int laneid) {
|
||||
row_reduce<base_ops::max, RV, RT, true>(row_accum, src, row_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the minimum of each row of the src register tile in the row_accum column vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_min(thread RV &row_accum, thread const RT &src, const int laneid) {
|
||||
row_reduce<base_ops::min, RV, RT, true>(row_accum, src, row_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the sum of each row of the src register tile in the row_accum column vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_sum(thread RV &row_accum, thread const RT &src, const int laneid) {
|
||||
row_reduce<base_ops::sum, RV, RT, true>(row_accum, src, row_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the product of each row of the src register tile in the row_accum column vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_prod(thread RV &row_accum, thread const RT &src, const int laneid) {
|
||||
row_reduce<base_ops::mul, RV, RT, true>(row_accum, src, row_accum, laneid);
|
||||
}
|
||||
|
||||
// three-operand row reductions. (Accumulate ONTO.)
|
||||
/**
|
||||
* @brief Store the maximum of each row of the src register tile, as well as the src_accum column vector, in the row_accum column vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_max(thread RV &row_accum, thread const RT &src, thread const RV &src_accum, const int laneid) {
|
||||
// using T = typename RV::T;
|
||||
// using T2 = typename RV::T2;
|
||||
// const short leader = (laneid / 16) * 16 + ((laneid / 2) % 4) * 2;
|
||||
//
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < src.height; i++) {
|
||||
// T accum_thread = metal::max(src.tiles[i][0].data.thread_elements()[0], src.tiles[i][0].data.thread_elements()[1]);
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 1; j < src.width; j++) {
|
||||
// accum_thread = metal::max(accum_thread, src.tiles[i][j].data.thread_elements()[0]);
|
||||
// accum_thread = metal::max(accum_thread, src.tiles[i][j].data.thread_elements()[1]);
|
||||
// }
|
||||
// accum_thread = metal::max(accum_thread, shfl_down_sync<T>(accum_thread, 1));
|
||||
// accum_thread = metal::max(accum_thread, shfl_down_sync<T>(accum_thread, 8));
|
||||
// accum_thread = shfl_sync<T>(accum_thread, leader);
|
||||
// if(false) { row_accum[i][0] = accum_thread; }
|
||||
// else { row_accum[i][0] = metal::max(src_accum[i][0], accum_thread); }
|
||||
// }
|
||||
|
||||
row_reduce<base_ops::max, RV, RT, false>(row_accum, src, src_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the minimum of each row of the src register tile, as well as the src_accum column vector, in the row_accum column vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_min(thread RV &row_accum, thread const RT &src, thread const RV &src_accum, const int laneid) {
|
||||
row_reduce<base_ops::min, RV, RT, false>(row_accum, src, src_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the sum of each row of the src register tile, as well as the src_accum column vector, in the row_accum column vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_sum(thread RV &row_accum, thread const RT &src, thread const RV &src_accum, const int laneid) {
|
||||
// using T = typename RV::T;
|
||||
// using T2 = typename RV::T2;
|
||||
// const short leader = (laneid / 16) * 16 + ((laneid / 2) % 4) * 2;
|
||||
//
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < src.height; i++) {
|
||||
// T accum_thread = (src.tiles[i][0].data.thread_elements()[0] + src.tiles[i][0].data.thread_elements()[1]);
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 1; j < src.width; j++) {
|
||||
// accum_thread = (accum_thread + src.tiles[i][j].data.thread_elements()[0]);
|
||||
// accum_thread = (accum_thread + src.tiles[i][j].data.thread_elements()[1]);
|
||||
// }
|
||||
// T shfl_val = shfl_down_sync<T>(accum_thread, 1);
|
||||
// accum_thread = (accum_thread + shfl_val);
|
||||
// shfl_val = shfl_down_sync<T>(accum_thread, 8);
|
||||
// accum_thread = (accum_thread + shfl_val);
|
||||
// accum_thread = shfl_sync<T>(accum_thread, leader);
|
||||
//// accum_thread = metal::simd_sum(accum_thread);
|
||||
// if(false) {
|
||||
// row_accum[i][0] = accum_thread;
|
||||
// }
|
||||
// else {
|
||||
// T src_val = src_accum[i][0];
|
||||
// row_accum[i][0] = (src_val + accum_thread);
|
||||
// }
|
||||
// }
|
||||
row_reduce<base_ops::sum, RV, RT, false>(row_accum, src, src_accum, laneid);
|
||||
}
|
||||
//template<typename RV, typename RT>
|
||||
//static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
//row_sum(thread RV &row_accum, thread const RT &src, thread const RV &src_accum, const int laneid, const int warpId, threadgroup typename RT::T* smem) {
|
||||
// using T = typename RV::T;
|
||||
// using T2 = typename RV::T2;
|
||||
// using T4 = typename base_types::packing<T>::packed_four;
|
||||
// const short leader = (laneid / 16) * 16 + ((laneid / 2) % 4) * 2;
|
||||
// const short qid = laneid / 4;
|
||||
// const int offsetX = (qid & 4) + (laneid / 2) % 4;
|
||||
// const int offsetY = (qid & 2) + laneid % 2;
|
||||
// const int smem_idx_row = 32 * warpId + offsetY * 4;
|
||||
// const int smem_idx = smem_idx_row + offsetX;
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int i = 0; i < src.height; i++) {
|
||||
// T accum_thread = src.tiles[i][0].data.thread_elements()[0] + src.tiles[i][0].data.thread_elements()[1];
|
||||
// #pragma clang loop unroll(full)
|
||||
// for(int j = 1; j < src.width; j++) {
|
||||
// accum_thread = accum_thread + src.tiles[i][0].data.thread_elements()[0];
|
||||
// accum_thread = accum_thread + src.tiles[i][0].data.thread_elements()[1];
|
||||
// }
|
||||
// {
|
||||
// metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
// smem[smem_idx] = accum_thread;
|
||||
// metal::simdgroup_barrier(metal::mem_flags::mem_threadgroup);
|
||||
// T4 vals = *(threadgroup T4*)(&smem[smem_idx_row]);
|
||||
// accum_thread = vals[0] + vals[1] + vals[2] + vals[3];
|
||||
// }
|
||||
// row_accum[i][0] = src_accum[i][0] + accum_thread;
|
||||
//
|
||||
// }
|
||||
//}
|
||||
|
||||
|
||||
/**
|
||||
* @brief Store the product of each row of the src register tile, as well as the src_accum column vector, in the row_accum column vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
row_prod(thread RV &row_accum, thread const RT &src, thread const RV &src_accum, const int laneid) {
|
||||
row_reduce<base_ops::mul, RV, RT, false>(row_accum, src, src_accum, laneid);
|
||||
}
|
||||
// two-operand col reductions. (Accumulate and REPLACE.)
|
||||
|
||||
/**
|
||||
* @brief Store the maximum of each column of the src register tile in the col_accum row vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_max(thread RV &col_accum, thread const RT &src, const int laneid) {
|
||||
col_reduce<base_ops::max, RV, RT, true>(col_accum, src, col_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the minimum of each column of the src register tile in the col_accum row vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_min(thread RV &col_accum, thread const RT &src, const int laneid) {
|
||||
col_reduce<base_ops::min, RV, RT, true>(col_accum, src, col_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the sum of each column of the src register tile in the col_accum row vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_sum(thread RV &col_accum, thread const RT &src, const int laneid) {
|
||||
col_reduce<base_ops::sum, RV, RT, true>(col_accum, src, col_accum, laneid);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store the product of each column of the src register tile in the col_accum row vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_prod(thread RV &col_accum, thread const RT &src, const int laneid) {
|
||||
col_reduce<base_ops::mul, RV, RT, true>(col_accum, src, col_accum, laneid);
|
||||
}
|
||||
// three-operand col reductions. (Accumulate ONTO.)
|
||||
/**
|
||||
* @brief Store the maximum of each column of the src register tile, as well as the src_accum row vector, in the col_accum row vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_max(thread RV &col_accum, thread const RT &src, thread const RV &src_accum, const int laneid) {
|
||||
col_reduce<base_ops::max, RV, RT, false>(col_accum, src, src_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the minimum of each column of the src register tile, as well as the src_accum row vector, in the col_accum row vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_min(thread RV &col_accum, thread const RT &src, thread const RV &src_accum, const int laneid) {
|
||||
col_reduce<base_ops::min, RV, RT, false>(col_accum, src, src_accum, laneid);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store the sum of each column of the src register tile, as well as the src_accum row vector, in the col_accum row vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_sum(thread RV &col_accum, thread const RT &src, thread const RV &src_accum, const int laneid) {
|
||||
col_reduce<base_ops::sum, RV, RT, false>(col_accum, src, src_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the product of each column of the src register tile, as well as the src_accum row vector, in the col_accum row vector.
|
||||
*
|
||||
* @tparam V The vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename RV, typename RT>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_tile<RT>() && ducks::is_register_vector<RV>(), void>::type
|
||||
col_prod(thread RV &col_accum, thread const RT &src, thread const RV &src_accum, const int laneid) {
|
||||
col_reduce<base_ops::mul, RV, RT, false>(col_accum, src, src_accum, laneid);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief An aggregate header for warp operations on register tiles.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "conversions.metal"
|
||||
#include "maps.metal"
|
||||
#include "mma.metal"
|
||||
#include "reductions.metal"
|
||||
@@ -0,0 +1,162 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Conversions on vectors stored in registers.
|
||||
*/
|
||||
|
||||
#pragma once // done
|
||||
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
|
||||
namespace detail {
|
||||
static METAL_FUNC int colstart_from_laneid(const int laneid) { // rowvec
|
||||
return (laneid % 2) * 2 + ((laneid / 8) % 2) * 4;
|
||||
}
|
||||
// 0,1,2,3,4,5,6,7 -> 0,2,1,3,8,10,9,11
|
||||
static METAL_FUNC int leader_from_col(const int col) { // rowvec
|
||||
return (col / 4) * 8 + (col / 2) % 2 + (col % 2) * 2;
|
||||
}
|
||||
// 0,2,1,3,8,10,9,11 -> 0,1,0,1,0,1,0,1
|
||||
static METAL_FUNC int idx_from_colleader(const int laneid) { // rowvec
|
||||
return ((laneid % 8) / 2) % 2; // % 2 to protect against non-leaders
|
||||
}
|
||||
|
||||
static METAL_FUNC int row_from_laneid(const int laneid) { // rowvec
|
||||
return (laneid / 2) % 4 + (laneid / 16) * 4;
|
||||
}
|
||||
// 0,1,2,3,4,5,6,7 -> 0, 2, 4, 6, 16, 18, 20, 22
|
||||
static METAL_FUNC int leader_from_row(const int row) { // rowvec
|
||||
return (row/4) * 16 + (row % 4) * 2;
|
||||
}
|
||||
|
||||
|
||||
/* ----- ducks::is_align_register_vector<RV1>() && ducks::is_naive_register_vector<RV2>() -----*/
|
||||
static METAL_FUNC int col_leader_from_naive_laneid(const int laneid) { // rowvec
|
||||
int tile_col = laneid % 8;
|
||||
int base_leader = (tile_col / 4) * 8 + (tile_col / 2) % 2 + (tile_col % 2) * 16;
|
||||
return base_leader + 2 * (laneid / 8);
|
||||
}
|
||||
|
||||
static METAL_FUNC int local_send_idx_from_col(const int laneid) {
|
||||
return laneid >= 16;
|
||||
}
|
||||
|
||||
static METAL_FUNC int src_basetile_from_laneid(const int laneid) { // rowvec
|
||||
return (laneid/ 2) % 4;
|
||||
}
|
||||
|
||||
/* ----- ducks::is_ortho_register_vector<RV1>() && ducks::is_naive_register_vector<RV2>() -----*/
|
||||
static METAL_FUNC int row_leader_from_naive_laneid(const int laneid) { // rowvec
|
||||
int row = laneid % 8;
|
||||
int base_row = (row/4) * 16 + (row % 4) * 2;
|
||||
return base_row + (laneid / 8) % 2 + (laneid >= 16) * 8;
|
||||
}
|
||||
|
||||
static METAL_FUNC int ortho_send_tile_from_laneid(const int laneid) { // rowvec
|
||||
// uint32_t MASK_1 = 0b00000000010101010000000001010101;
|
||||
uint32_t MASK_2 = 0b00000000101010100000000010101010;
|
||||
uint32_t MASK_3 = 0b01010101000000000101010100000000;
|
||||
uint32_t MASK_4 = 0b10101010000000001010101000000000;
|
||||
return ((MASK_2 >> laneid) & 1) + ((MASK_3 >> laneid) & 1) * 2 + ((MASK_4 >> laneid) & 1) * 3;
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
}
|
||||
/**
|
||||
* @brief Copies data from one register vector to another.
|
||||
*
|
||||
* @tparam RV1 The type of the destination register vector.
|
||||
* @tparam RV2 The type of the source register vector.
|
||||
* @param dst[out] The destination register vector.
|
||||
* @param src[in] The source register vector to copy from.
|
||||
*/
|
||||
template<typename RV2, typename RV1>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV1>() && ducks::is_register_vector<RV2>(), void>::type
|
||||
copy(thread RV2 &dst, thread const RV1 &src, const ushort laneid) {
|
||||
static_assert(RV1::length == RV2::length, "Outer dimensions of the register vectors must be the same.");
|
||||
using D1 = typename RV1::dtype;
|
||||
using D2 = typename RV2::dtype;
|
||||
if (metal::is_same_v<typename RV1::layout, typename RV2::layout>) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < RV1::outer_dim; i++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < RV1::inner_dim; j++) {
|
||||
dst[i][j] = base_types::convertor<D1, D2>::convert(src[i][j]);
|
||||
}
|
||||
}
|
||||
} else if (ducks::is_align_register_vector<RV1>() && ducks::is_ortho_register_vector<RV2>()) { // align vector -> ortho vector
|
||||
const int row = detail::row_from_laneid(laneid);
|
||||
const int laneid_src = detail::leader_from_col(row);
|
||||
const int send_idx = detail::idx_from_colleader(laneid);
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < RV1::outer_dim; i++) {
|
||||
dst[i][0] = base_types::convertor<D1,D2>::convert(shfl_sync<D2>(src[i][send_idx], laneid_src));
|
||||
// dst[i][0] = 1;
|
||||
}
|
||||
} else if (ducks::is_ortho_register_vector<RV1>() && ducks::is_align_register_vector<RV2>()) { // ortho vector -> align vector
|
||||
const int col1 = detail::colstart_from_laneid(laneid);
|
||||
const int col2 = col1 + 1;
|
||||
const int laneid_src1 = detail::leader_from_row(col1);
|
||||
const int laneid_src2 = detail::leader_from_row(col2);
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < RV1::outer_dim; i++) {
|
||||
dst[i][0] = base_types::convertor<D2,D1>::convert(shfl_sync<D1>(src[i][0], laneid_src1));
|
||||
dst[i][1] = base_types::convertor<D2,D1>::convert(shfl_sync<D1>(src[i][0], laneid_src2));
|
||||
}
|
||||
} else if (ducks::is_align_register_vector<RV1>() && ducks::is_naive_register_vector<RV2>()) {
|
||||
const int src_laneid = detail::col_leader_from_naive_laneid(laneid);
|
||||
int align_send_tile = detail::src_basetile_from_laneid(laneid);
|
||||
int align_local_send_idx = detail::local_send_idx_from_col(laneid);
|
||||
int naive_tile_idx = 0;
|
||||
for (int l_idx = 0;
|
||||
l_idx < RV2::length;
|
||||
l_idx += 32, naive_tile_idx++, align_send_tile += 4)
|
||||
{
|
||||
D1 send_val = 0;
|
||||
if (align_send_tile < RV1::outer_dim) send_val = src[align_send_tile][align_local_send_idx];
|
||||
D1 recieve_val = shfl_sync<D1>(send_val, src_laneid);
|
||||
if (l_idx + laneid < RV2::length) dst[l_idx / 32][0] = base_types::convertor<D2,D1>::convert(recieve_val);
|
||||
}
|
||||
} else if (ducks::is_naive_register_vector<RV1>() && ducks::is_align_register_vector<RV2>()) {
|
||||
int col1 = detail::colstart_from_laneid(laneid);
|
||||
int col2 = col1 + 1;
|
||||
for (int i = 0; i < RV2::outer_dim; i++) {
|
||||
int src1 = (i%4) * 8 + col1;
|
||||
int src2 = (i%4) * 8 + col2;
|
||||
D1 send_val = src[i / 4][0];
|
||||
D1 recieve_val1 = shfl_sync<D1>(send_val, src1);
|
||||
D1 recieve_val2 = shfl_sync<D1>(send_val, src2);
|
||||
dst[i][0] = recieve_val1;
|
||||
dst[i][1] = recieve_val2;
|
||||
}
|
||||
} else if (ducks::is_ortho_register_vector<RV1>() && ducks::is_naive_register_vector<RV2>()) {
|
||||
const int src_laneid = detail::row_leader_from_naive_laneid(laneid);
|
||||
int ortho_send_tile = detail::ortho_send_tile_from_laneid(laneid);
|
||||
int naive_tile_idx = 0;
|
||||
for (int l_idx = 0; l_idx < RV2::length;
|
||||
l_idx += 32, naive_tile_idx++, ortho_send_tile += 4)
|
||||
{
|
||||
D1 send_val = 10;
|
||||
if (ortho_send_tile < RV1::outer_dim) send_val = src[ortho_send_tile][0];
|
||||
D1 recieve_val = shfl_sync<D1>(send_val, src_laneid);
|
||||
if (l_idx + laneid < RV2::length) dst[l_idx / 32][0] = base_types::convertor<D2,D1>::convert(recieve_val);
|
||||
}
|
||||
} else if (ducks::is_naive_register_vector<RV1>() && ducks::is_ortho_register_vector<RV2>()) {
|
||||
int row = detail::row_from_laneid(laneid);
|
||||
for (int i = 0; i < RV2::outer_dim; i++) {
|
||||
int src_laneid = (i%4) * 8 + row;
|
||||
D1 send_val = src[i / 4][0];
|
||||
D1 recieve_val = shfl_sync<D1>(send_val, src_laneid);
|
||||
dst[i][0] = recieve_val;
|
||||
}
|
||||
}
|
||||
else {
|
||||
// static_assert(RV1::inner_dim == RV2::inner_dim, "Something has gone deeply wrong with how register vectors were instantiated.");
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,288 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Maps on vectors stored in registers.
|
||||
*/
|
||||
|
||||
#pragma once // doneington
|
||||
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
|
||||
/* ---------- Vector Maps ---------- */
|
||||
|
||||
/**
|
||||
* @brief Perform a unary operation on a vector.
|
||||
*
|
||||
* @tparam op The unary operation to perform.
|
||||
* @tparam T The type of the vector.
|
||||
* @param dst[out] The destination vector where the result is stored.
|
||||
* @param src[in] The source vector to perform the operation on.
|
||||
*/
|
||||
template<typename op, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
unary_op(thread RV &dst, thread const RV &src) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.outer_dim; i++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.inner_dim; j++) {
|
||||
dst[i][j] = op::template op<typename RV::dtype>(src[i][j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
/**
|
||||
* @brief Perform a binary operation on two vectors.
|
||||
*
|
||||
* @tparam op The binary operation to perform.
|
||||
* @tparam T The type of the vectors.
|
||||
* @param dst[out] The destination vector where the result is stored.
|
||||
* @param lhs[in] The left-hand side vector for the operation.
|
||||
* @param rhs[in] The right-hand side vector for the operation.
|
||||
*/
|
||||
template<typename op, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
bin_op(thread RV &dst, thread const RV &lhs, thread const RV &rhs) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.outer_dim; i++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.inner_dim; j++) {
|
||||
dst[i][j] = op::template op<typename RV::dtype>(lhs[i][j], rhs[i][j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
/**
|
||||
* @brief Perform a binary operation on a vector and a scalar.
|
||||
*
|
||||
* @tparam op The binary operation to perform.
|
||||
* @tparam T The type of the vector.
|
||||
* @param dst[out] The destination vector where the result is stored.
|
||||
* @param src[in] The source vector for the operation.
|
||||
* @param param[in] The scalar parameter for the operation.
|
||||
*/
|
||||
template<typename op, typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
bin_op(thread RV &dst, thread const RV &src, thread const typename RV::dtype ¶m) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 0; i < dst.outer_dim; i++) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int j = 0; j < dst.inner_dim; j++) {
|
||||
dst[i][j] = op::template op<typename RV::dtype>(src[i][j], param);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
|
||||
|
||||
// ---- const ops ----
|
||||
|
||||
/**
|
||||
* @brief Sets all elements of a register vector to zero.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @param dst[out] Destination vector to be set to zero.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
zero(thread RV &dst) {
|
||||
unary_op<base_ops::zero, RV>(dst, dst);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Sets all elements of a register vector to one.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @param dst[out] Destination vector to be set to one.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
one(thread RV &dst) {
|
||||
unary_op<base_ops::one, RV>(dst, dst);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of a register vector to positive infinity.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @param dst[out] Destination vector to be set to positive infinity.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
pos_infty(thread RV &dst) {
|
||||
unary_op<base_ops::pos_infty, RV>(dst, dst);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of a register vector to negative infinity.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @param dst[out] Destination vector to be set to negative infinity.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
neg_infty(thread RV &dst) {
|
||||
unary_op<base_ops::neg_infty, RV>(dst, dst);
|
||||
}
|
||||
|
||||
// ---- unary ops ----
|
||||
|
||||
/**
|
||||
* @brief Copies the elements from one register vector to another.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @tparam U Type of the source vector.
|
||||
* @param dst[out] Destination vector where the elements will be copied to.
|
||||
* @param src[in] Source vector to copy the elements from.
|
||||
*/
|
||||
template<typename RV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>() && ducks::base_types::isT1Type<U>(), void>::type
|
||||
copy(thread RV &dst, thread const U &src) {
|
||||
bin_op<base_ops::copy2, RV>(dst, dst, src); // the second arg is ignored here.
|
||||
}
|
||||
/**
|
||||
* @brief Applies the exponential function element-wise to a register vector.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @param dst[out] Destination vector where the exponential values will be stored.
|
||||
* @param src[in] Source vector to apply the exponential function to.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
exp(thread RV &dst, thread const RV &src) {
|
||||
unary_op<base_ops::exp, RV>(dst, src);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the exponential function element-wise to a register vector, in base 2.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @param dst[out] Destination vector where the exponential values will be stored.
|
||||
* @param src[in] Source vector to apply the exponential function to.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
exp2(thread RV &dst, thread const RV &src) {
|
||||
unary_op<base_ops::exp2, RV>(dst, src);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the natural logarithm function element-wise to a register vector.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @param dst[out] Destination vector where the exponential values will be stored.
|
||||
* @param src[in] Source vector to apply the exponential function to.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
log(thread RV &dst, thread const RV &src) {
|
||||
unary_op<base_ops::log, RV>(dst, src);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the absolute value function element-wise to a register vector.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @param dst[out] Destination vector where the absolute values will be stored.
|
||||
* @param src[in] Source vector to apply the absolute value function to.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
abs(thread RV &dst, thread const RV &src) {
|
||||
unary_op<base_ops::abs, RV>(dst, src);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the rectified linear unit (ReLU) function element-wise to a register vector.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @param dst[out] Destination vector where the ReLU values will be stored.
|
||||
* @param src[in] Source vector to apply the ReLU function to.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
relu(thread RV &dst, thread const RV &src) {
|
||||
unary_op<base_ops::relu, RV>(dst, src);
|
||||
}
|
||||
|
||||
// ---- binary ops ----
|
||||
|
||||
/**
|
||||
* @brief Computes the element-wise maximum of two register vectors.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the maximum values will be stored.
|
||||
* @param lhs[in] First vector for the maximum operation.
|
||||
* @param rhs[in] Second vector for the maximum operation.
|
||||
*/
|
||||
template<typename RV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
max(thread RV &dst, thread const RV &lhs, thread const U &rhs) {
|
||||
bin_op<base_ops::max, RV>(dst, lhs, rhs);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise minimum of two register vectors.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the minimum values will be stored.
|
||||
* @param lhs[in] First vector for the minimum operation.
|
||||
* @param rhs[in] Second vector for the minimum operation.
|
||||
*/
|
||||
template<typename RV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
min(thread RV &dst, thread const RV &lhs, thread const U &rhs) {
|
||||
bin_op<base_ops::min, RV>(dst, lhs, rhs);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise sum of two register vectors.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the sum values will be stored.
|
||||
* @param lhs[in] First vector for the sum operation.
|
||||
* @param rhs[in] Second vector for the sum operation.
|
||||
*/
|
||||
template<typename RV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
add(thread RV &dst, thread const RV &lhs, thread const U &rhs) {
|
||||
bin_op<base_ops::sum, RV>(dst, lhs, rhs);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise difference of two register vectors.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the difference values will be stored.
|
||||
* @param lhs[in] First vector for the difference operation.
|
||||
* @param rhs[in] Second vector for the difference operation.
|
||||
*/
|
||||
template<typename RV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
sub(thread RV &dst, thread const RV &lhs, thread const U &rhs) {
|
||||
bin_op<base_ops::sub, RV>(dst, lhs, rhs);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise product of two register vectors.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the product values will be stored.
|
||||
* @param lhs[in] First vector for the product operation.
|
||||
* @param rhs[in] Second vector for the product operation.
|
||||
*/
|
||||
template<typename RV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
mul(thread RV &dst, thread const RV &lhs, thread const U &rhs) {
|
||||
bin_op<base_ops::mul, RV>(dst, lhs, rhs);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise division of two register vectors.
|
||||
*
|
||||
* @tparam T Register vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the division values will be stored.
|
||||
* @param lhs[in] First vector for the division operation.
|
||||
* @param rhs[in] Second vector for the division operation.
|
||||
*/
|
||||
template<typename RV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
div(thread RV &dst, thread const RV &lhs, thread const U &rhs) {
|
||||
bin_op<base_ops::div, RV>(dst, lhs, rhs);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Reductions on vectors stored in registers.
|
||||
*/
|
||||
|
||||
#pragma once // done
|
||||
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
/* ---------- Vector Reductions ---------- */
|
||||
|
||||
/**
|
||||
* @brief Performs a reduction operation on elements of a register vector within a warp.
|
||||
*
|
||||
* This function applies a specified operation to reduce the elements of a register vector `src` to a single value.
|
||||
* The result is stored in `accum`. If the `reset` parameter is true, the reduction includes an initial value `src_accum`.
|
||||
* The reduction operation is performed in a warp-wide context, ensuring synchronization between threads in the warp.
|
||||
*
|
||||
* @tparam op The operation to perform on the elements. Must provide a static `op` method.
|
||||
* @tparam RV The type of the register vector. Must satisfy the `ducks::rv::all` concept.
|
||||
* @tparam reset A boolean flag indicating whether to include an initial value in the reduction.
|
||||
* @param[out] accum The result of the reduction operation.
|
||||
* @param[in] src The register vector to reduce.
|
||||
* @param[in] src_accum The initial value to include in the reduction if `reset` is false.
|
||||
*/
|
||||
template<typename op, typename RV, bool reset>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
reduce(
|
||||
thread typename RV::T &dst_accum,
|
||||
thread const RV &src,
|
||||
thread const typename RV::T &src_accum,
|
||||
const ushort laneid) {
|
||||
using T = typename RV::T;
|
||||
if (ducks::is_ortho_register_vector<RV>()) { // col vector
|
||||
T accum = src[0][0];
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 1; i < src.outer_dim; i++) {
|
||||
accum = op::template op<T>(accum, src[i][0]);
|
||||
}
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 2));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 4));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 16));
|
||||
if (!reset) accum = op::template op<T>(accum, src_accum);
|
||||
dst_accum = shfl_sync(accum, 0);
|
||||
}
|
||||
else if (ducks::is_align_register_vector<RV>()) { // row vector
|
||||
T accum = op::template op<T>(src[0][0], src[0][1]);
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 1; i < src.outer_dim; i++) {
|
||||
accum = op::template op<T>(accum, src[i][0]);
|
||||
accum = op::template op<T>(accum, src[i][1]);
|
||||
}
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 1));
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 8));
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
|
||||
accum = shfl_sync<T>(accum, 0);
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
if (!reset) accum = op::template op<T>(accum, src_accum);
|
||||
dst_accum = accum;
|
||||
}
|
||||
else if (ducks::is_naive_register_vector<RV>()) {
|
||||
// T accum = src[0][0];
|
||||
T accum;
|
||||
if (laneid < src.length) accum = src[0][0];
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = 1; i < src.outer_dim; i++) {
|
||||
if (i*SIMD_THREADS + laneid < src.length) {
|
||||
accum = op::template op<T>(accum, src[i][0]);
|
||||
}
|
||||
}
|
||||
if (src.length == 8) {
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 1));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 2));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 4));
|
||||
} else if (src.length == 16) {
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 1));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 2));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 4));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 8));
|
||||
} else if (src.length == 24) {
|
||||
if (laneid < 24) {
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 1));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 2));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 4));
|
||||
|
||||
T shfle_val = shfl_down_sync<T>(accum, 8);
|
||||
if (laneid < 16) {
|
||||
accum = op::template op<T>(accum, shfle_val);
|
||||
}
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 16));
|
||||
}
|
||||
|
||||
} else {
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 1));
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 2));
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 4));
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 8));
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 16));
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
}
|
||||
|
||||
if (!reset) accum = op::template op<T>(accum, src_accum);
|
||||
dst_accum = shfl_sync(accum, 0);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Finds the maximum element in a register vector.
|
||||
*
|
||||
* @tparam RV The type of the register vector. Must satisfy the `ducks::rv::all` concept.
|
||||
* @param[out] max_val The maximum value found in the vector.
|
||||
* @param[in] src The register vector to find the maximum in.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
max(thread typename base_types::packing<typename RV::dtype>::unpacked_type &max_val, thread const RV &src, const ushort laneid) {
|
||||
reduce<base_ops::max, RV, true>(max_val, src, max_val, laneid);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Finds the minimum element in a register vector.
|
||||
*
|
||||
* @tparam RV The type of the register vector. Must satisfy the `ducks::rv::all` concept.
|
||||
* @param[out] min_val The minimum value found in the vector.
|
||||
* @param[in] src The register vector to find the minimum in.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
min(thread typename base_types::packing<typename RV::dtype>::unpacked_type &min_val, thread const RV &src, const ushort laneid) {
|
||||
reduce<base_ops::min, RV, true>(min_val, src, min_val, laneid);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Calculates the sum of elements in a register vector.
|
||||
*
|
||||
* @tparam RV The type of the register vector. Must satisfy the `ducks::rv::all` concept.
|
||||
* @param[out] sum_val The sum of the values in the vector.
|
||||
* @param[in] src The register vector to sum.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
sum(thread typename base_types::packing<typename RV::dtype>::unpacked_type &sum_val, thread const RV &src, const ushort laneid) {
|
||||
reduce<base_ops::sum, RV, true>(sum_val, src, sum_val, laneid);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Calculates the product of elements in a register vector.
|
||||
*
|
||||
* @tparam RV The type of the register vector. Must satisfy the `ducks::rv::all` concept.
|
||||
* @param[out] prod_val The product of the values in the vector.
|
||||
* @param[in] src The register vector to multiply.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
prod(thread typename base_types::packing<typename RV::dtype>::unpacked_type &prod_val, thread const RV &src, const ushort laneid) {
|
||||
reduce<base_ops::mul, RV, true>(prod_val, src, prod_val, laneid);
|
||||
}
|
||||
|
||||
// Three operand versions.
|
||||
|
||||
/**
|
||||
* @brief Finds the maximum element in a register vector and accumulates it with src_accum.
|
||||
*
|
||||
* @tparam RV The type of the register vector. Must satisfy the `ducks::rv::all` concept.
|
||||
* @param[out] max_val The maximum value found in the vector, accumulated with src_accum.
|
||||
* @param[in] src The register vector to find the maximum in.
|
||||
* @param[in] src_accum The initial value to accumulate with the maximum value found.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
max(thread typename base_types::packing<typename RV::dtype>::unpacked_type &max_val,
|
||||
thread const RV &src,
|
||||
thread const typename base_types::packing<typename RV::dtype>::unpacked_type &src_accum, const ushort laneid) {
|
||||
reduce<base_ops::max, RV, false>(max_val, src, src_accum, laneid);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Finds the minimum element in a register vector and accumulates it with src_accum.
|
||||
*
|
||||
* @tparam RV The type of the register vector. Must satisfy the `ducks::rv::all` concept.
|
||||
* @param[out] min_val The minimum value found in the vector, accumulated with src_accum.
|
||||
* @param[in] src The register vector to find the minimum in.
|
||||
* @param[in] src_accum The initial value to accumulate with the minimum value found.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
min(thread typename base_types::packing<typename RV::dtype>::unpacked_type &min_val,
|
||||
thread const RV &src,
|
||||
thread const typename base_types::packing<typename RV::dtype>::unpacked_type &src_accum, const ushort laneid) {
|
||||
reduce<base_ops::min, RV, false>(min_val, src, src_accum, laneid);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Calculates the sum of elements in a register vector and accumulates it with src_accum.
|
||||
*
|
||||
* @tparam RV The type of the register vector. Must satisfy the `ducks::rv::all` concept.
|
||||
* @param[out] sum_val The sum of the values in the vector, accumulated with src_accum.
|
||||
* @param[in] src The register vector to sum.
|
||||
* @param[in] src_accum The initial value to accumulate with the sum of the vector.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
sum(thread typename base_types::packing<typename RV::dtype>::unpacked_type &sum_val,
|
||||
thread const RV &src,
|
||||
thread const typename base_types::packing<typename RV::dtype>::unpacked_type &src_accum, const ushort laneid) {
|
||||
reduce<base_ops::sum, RV, false>(sum_val, src, src_accum, laneid);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Calculates the product of elements in a register vector and accumulates it with src_accum.
|
||||
*
|
||||
* @tparam RV The type of the register vector. Must satisfy the `ducks::rv::all` concept.
|
||||
* @param[out] prod_val The product of the values in the vector, accumulated with src_accum.
|
||||
* @param[in] src The register vector to multiply.
|
||||
* @param[in] src_accum The initial value to accumulate with the product of the vector.
|
||||
*/
|
||||
template<typename RV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
|
||||
prod(thread typename base_types::packing<typename RV::dtype>::unpacked_type &prod_val,
|
||||
thread const RV &src,
|
||||
thread const typename base_types::packing<typename RV::dtype>::unpacked_type &src_accum, const ushort laneid) {
|
||||
reduce<base_ops::mul, RV, false>(prod_val, src, src_accum, laneid);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
#pragma once
|
||||
#include "conversions.metal"
|
||||
#include "maps.metal"
|
||||
#include "reductions.metal"
|
||||
@@ -0,0 +1,3 @@
|
||||
#pragma once
|
||||
#include "tile/tile.metal"
|
||||
#include "vec/vec.metal"
|
||||
@@ -0,0 +1,59 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Conversions between shared tile types.
|
||||
*/
|
||||
|
||||
#pragma once // not done, add subtile
|
||||
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
/* ---------- COPIES ---------- */
|
||||
/**
|
||||
* @brief Copies data from one shared memory tile to another, potentially with different data types and layouts.
|
||||
*
|
||||
* @tparam T The data type of the destination tile.
|
||||
* @tparam U The data type of the source tile.
|
||||
* @tparam _height The height of the tile.
|
||||
* @tparam _width The width of the tile.
|
||||
* @tparam L1 The layout of the destination tile.
|
||||
* @tparam L2 The layout of the source tile.
|
||||
* @param[out] dst The destination tile.
|
||||
* @param[in] src The source tile.
|
||||
*/
|
||||
template<typename T, typename U, int _height, int _width>
|
||||
static METAL_FUNC void copy(threadgroup st<T, _height, _width> &dst, threadgroup const st<U, _height, _width> &src, const ushort laneid) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid; i < dst.num_elements; i+=mittens::SIMD_THREADS) {
|
||||
int row = i/dst.cols, col = i%dst.cols;
|
||||
dst[{row, col}] = base_types::convertor<T, U>::convert(src[{row, col}]);
|
||||
}
|
||||
}
|
||||
|
||||
///* ---------- SUBTILE ---------- */
|
||||
//
|
||||
///**
|
||||
//* @brief Returns a reference to a subtile of the given shared tile.
|
||||
//*
|
||||
//* @tparam subtile_height The height of the subtile.
|
||||
//* @tparam subtile_width The width of the subtile.
|
||||
//* @tparam ST The type of the input tile, which must satisfy the ducks::st::all concept.
|
||||
//* @param src The input tile.
|
||||
//* @param row_idx The row index of the subtile, in units of subtile_height*16 elements.
|
||||
//* @param col_idx The col index of the subtile, in units of subtile_width*16 elements.
|
||||
//* @return A reference to the subtile.
|
||||
//*
|
||||
//* @note The subtile {height, width} must evenly divide the tile {height, width}.
|
||||
//*/
|
||||
//template<int subtile_height, int subtile_width, ducks::st::all ST>
|
||||
//__device__ inline typename ST::subtile<subtile_height, subtile_width> subtile_inplace(ST &src, int row_idx, int col_idx) {
|
||||
// static_assert(ST::height % subtile_height == 0);
|
||||
// static_assert(ST::width % subtile_width == 0);
|
||||
// return typename ST::subtile<subtile_height, subtile_width>(
|
||||
// &src[0], subtile_height*16*row_idx, subtile_width*16*col_idx
|
||||
// );
|
||||
//}
|
||||
|
||||
}
|
||||
|
||||
@@ -0,0 +1,485 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Warp-scope maps on shared tiles.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
/* ---------- Uniform tile maps (independent of layout) ---------- */
|
||||
|
||||
/**
|
||||
* @brief Performs a uniform unary operation on a tile.
|
||||
*
|
||||
* This function applies a given unary operation to each element of the source tile and stores the result in the destination tile.
|
||||
* The operation is applied independently to each element, without considering its position or the values of neighboring elements.
|
||||
*
|
||||
* @tparam op The unary operation to be applied. Must be specialized to support operation on the data type of T.
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the unary operation is applied.
|
||||
*/
|
||||
template<typename op, typename ST> // T2, w, h can be inferred from dst as long as op is specialized
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
unary_map(threadgroup ST &dst, threadgroup const ST &src, const ushort laneid) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid; i < ST::num_elements; i += SIMD_THREADS) {
|
||||
dst.data[i] = op::template op<typename ST::dtype>(src.data[i]);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* @brief Performs a uniform binary operation on a tile with a scalar parameter.
|
||||
*
|
||||
* This function applies a given binary operation to each element of the source tile and a scalar parameter, then stores the result in the destination tile.
|
||||
* The operation is applied independently to each element, treating the scalar parameter as the second operand for each operation.
|
||||
*
|
||||
* @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T and the scalar parameter.
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the binary operation is applied.
|
||||
* @param[in] param The scalar parameter to be used as the second operand in the binary operation.
|
||||
*/
|
||||
template<typename op, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
bin_map(threadgroup ST &dst, threadgroup const ST &src, thread const typename ST::dtype ¶m, const short laneid) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid; i < dst.num_elements; i += SIMD_THREADS) {
|
||||
dst.data[i] = op::template op<typename ST::dtype>(src.data[i], param);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Performs a uniform binary operation on two tiles.
|
||||
*
|
||||
* This function applies a given binary operation to corresponding elements of two source tiles and stores the result in the destination tile.
|
||||
* The operation is applied independently to each pair of elements, without considering their positions or the values of neighboring elements.
|
||||
*
|
||||
* @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T.
|
||||
* @tparam T The type of the tiles. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile to which the binary operation is applied.
|
||||
* @param[in] rhs The second source tile to which the binary operation is applied.
|
||||
*/
|
||||
template<typename op, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
bin_map(threadgroup ST &dst, threadgroup const ST &lhs, threadgroup const ST &rhs, const ushort laneid) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid; i < dst.num_elements; i += SIMD_THREADS) {
|
||||
int row = i/dst.cols, col = i%dst.cols;
|
||||
dst.data[i] = op::template op<typename ST::dtype>(lhs.data[i], rhs.data[i]);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Performs a row-wise binary operation on a tile with a vector.
|
||||
*
|
||||
* This function applies a given binary operation to each row of the source tile and the corresponding element of the source vector,
|
||||
* then stores the result in the destination tile. The operation is applied independently to each row, using the vector element as
|
||||
* the second operand for each element in the row.
|
||||
*
|
||||
* @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T and the vector elements.
|
||||
* @tparam T The type of the tiles. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam V The type of the vector. Must have the same data type as T.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the binary operation is applied.
|
||||
* @param[in] vec The source vector containing the second operand for each row operation.
|
||||
*/
|
||||
template<typename op, typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>, void>::type
|
||||
row_map(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &vec, const ushort laneid) {
|
||||
static_assert(metal::is_same<typename ST::dtype, typename SV::dtype>::value, "Tile and vector must have the same data type");
|
||||
static_assert(SV::length == ST::rows, "Vector length must match the number of rows in the tile");
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid; i < dst.num_elements; i += SIMD_THREADS) {
|
||||
int row = i/ST::cols, col = i%ST::cols;
|
||||
dst[{row, col}] = op::template op<typename ST::dtype>(src[{row, col}], vec[row]);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Performs a column-wise binary operation on a tile with a vector.
|
||||
*
|
||||
* This function applies a given binary operation to each column of the source tile and the corresponding element of the source vector,
|
||||
* then stores the result in the destination tile. The operation is applied independently to each column, using the vector element as
|
||||
* the second operand for each element in the column.
|
||||
*
|
||||
* @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T and the vector elements.
|
||||
* @tparam T The type of the tiles. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam V The type of the vector. Must have the same data type as T.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the binary operation is applied.
|
||||
* @param[in] vec The source vector containing the second operand for each column operation.
|
||||
*/
|
||||
template<typename op, typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_map(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &vec, const ushort laneid) {
|
||||
static_assert(metal::is_same<typename ST::dtype, typename SV::dtype>::value, "Tile and vector must have the same data type");
|
||||
static_assert(SV::length == ST::cols, "Vector length must match the number of columns in the tile");
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid; i < dst.num_elements; i += SIMD_THREADS) {
|
||||
int row = i/dst.cols, col = i%dst.cols;
|
||||
dst[{row, col}] = op::template op<typename ST::dtype>(src[{row, col}], vec[col]);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
|
||||
|
||||
// const maps
|
||||
/**
|
||||
* @brief Sets all elements of the destination tile to zero.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
zero(threadgroup ST &dst, const ushort laneid) {
|
||||
unary_map<base_ops::zero, ST>(dst, dst, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of the destination tile to one.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
one(threadgroup ST &dst, const ushort laneid) {
|
||||
unary_map<base_ops::one, ST>(dst, dst, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of the destination tile to positive infinity.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
pos_infty(threadgroup ST &dst, const ushort laneid) {
|
||||
unary_map<base_ops::pos_infty, ST>(dst, dst, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of the destination tile to negative infinity.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
neg_infty(threadgroup ST &dst, const ushort laneid) {
|
||||
unary_map<base_ops::neg_infty, ST>(dst, dst, laneid);
|
||||
}
|
||||
|
||||
// unary maps
|
||||
/**
|
||||
* @brief Applies the exponential function to each element of the source tile and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the exponential function is applied.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
exp(threadgroup ST &dst, threadgroup const ST &src, const ushort laneid) {
|
||||
unary_map<base_ops::exp, ST>(dst, src, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the exponential function to each element of the source tile and stores the result in the destination tile, in base 2.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the exponential function is applied.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
exp2(threadgroup ST &dst, threadgroup const ST &src, const ushort laneid) {
|
||||
unary_map<base_ops::exp2, ST>(dst, src, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the natural logarithm function to each element of the source tile and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the natural logarithm function is applied.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
log(threadgroup ST &dst, threadgroup const ST &src, const ushort laneid) {
|
||||
unary_map<base_ops::log, ST>(dst, src, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the absolute function to each element of the source tile and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the absolute function is applied.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
abs(threadgroup ST &dst, threadgroup const ST &src, const ushort laneid) {
|
||||
unary_map<base_ops::abs, ST>(dst, src, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the rectified linear unit function to each element of the source tile and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source tile to which the rectified linear unit function is applied.
|
||||
*/
|
||||
template<typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
relu(threadgroup ST &dst, const threadgroup ST &src, const ushort laneid) {
|
||||
unary_map<base_ops::relu, ST>(dst, src, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Copies the elements of the source tile to the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] src The source data to be copied.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
copy(threadgroup ST &dst, thread const U &src, const ushort laneid) {
|
||||
bin_map<base_ops::copy2, ST>(dst, dst, src, laneid);
|
||||
}
|
||||
|
||||
// uniform binary maps
|
||||
/**
|
||||
* @brief Finds the maximum of each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile.
|
||||
* @param[in] rhs The second source data.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
max(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const ushort laneid) {
|
||||
bin_map<base_ops::max, ST>(dst, lhs, rhs, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Finds the minimum of each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile.
|
||||
* @param[in] rhs The second source data.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
min(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const ushort laneid) {
|
||||
bin_map<base_ops::min, ST>(dst, lhs, rhs, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Adds each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile.
|
||||
* @param[in] rhs The second source data.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
add(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const ushort laneid) {
|
||||
bin_map<base_ops::sum, ST>(dst, lhs, rhs, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Subtracts each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile.
|
||||
* @param[in] rhs The second source data.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
sub(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const ushort laneid) {
|
||||
bin_map<base_ops::sub, ST>(dst, lhs, rhs, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Multiplies each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile.
|
||||
* @param[in] rhs The second source data.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
mul(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const ushort laneid) {
|
||||
bin_map<base_ops::mul, ST>(dst, lhs, rhs, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Divides each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
|
||||
*
|
||||
* @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
|
||||
* @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
|
||||
* @param[out] dst The destination tile where the results are stored.
|
||||
* @param[in] lhs The first source tile.
|
||||
* @param[in] rhs The second source data.
|
||||
*/
|
||||
template<typename ST, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
div(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const ushort laneid) {
|
||||
bin_map<base_ops::div, ST>(dst, lhs, rhs, laneid);
|
||||
}
|
||||
|
||||
// Row and col maps
|
||||
|
||||
/**
|
||||
* @brief Adds row values to each row of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the addition on.
|
||||
* @param row_values[in] Column vector containing values to add to each row.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
add_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const ushort laneid) {
|
||||
row_map<base_ops::sum, ST, SV>(dst, src, row_values, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Subtracts row values from each row of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the subtraction on.
|
||||
* @param row_values[in] Column vector containing values to subtract from each row.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
sub_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const ushort laneid) {
|
||||
row_map<base_ops::sub, ST, SV>(dst, src, row_values, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Multiplies each row of a tile by row values.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the multiplication on.
|
||||
* @param row_values[in] Column vector containing values to multiply each row by.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
mul_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const ushort laneid) {
|
||||
row_map<base_ops::mul, ST, SV>(dst, src, row_values, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Divides each row of a tile by row values.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the division on.
|
||||
* @param row_values[in] Column vector containing values to divide each row by.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_tile<SV>(), void>::type
|
||||
div_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const ushort laneid) {
|
||||
row_map<base_ops::div, ST, SV>(dst, src, row_values, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Broadcast a vector into into a tile's rows.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Column vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param row_values[in] Column vector containing values to broadcast into rows.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
broadcast_row(threadgroup ST &dst, threadgroup const SV &row_values, const ushort laneid) {
|
||||
row_map<base_ops::copy2, ST, SV>(dst, dst, row_values, laneid);
|
||||
}
|
||||
|
||||
|
||||
// col maps
|
||||
/**
|
||||
* @brief Adds column values to each column of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the addition on.
|
||||
* @param col_values[in] Row vector containing values to add to each column.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
add_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const ushort laneid) {
|
||||
col_map<base_ops::sum, ST, SV>(dst, src, col_values, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Subtracts column values from each column of a tile.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the subtraction on.
|
||||
* @param col_values[in] Row vector containing values to subtract from each column.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
sub_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const ushort laneid) {
|
||||
col_map<base_ops::sub, ST, SV>(dst, src, col_values, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Multiplies each column of a tile by column values.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the multiplication on.
|
||||
* @param col_values[in] Row vector containing values to multiply each column by.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
mul_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const ushort laneid) {
|
||||
col_map<base_ops::mul, ST, SV>(dst, src, col_values, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Divides each column of a tile by column values.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param src[in] Source tile to apply the division on.
|
||||
* @param col_values[in] Row vector containing values to divide each column by.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
|
||||
div_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const ushort laneid) {
|
||||
col_map<base_ops::div, ST, SV>(dst, src, col_values, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Broadcast a vector into into a tile's columns.
|
||||
*
|
||||
* @tparam T Tile type.
|
||||
* @tparam V Row vector type.
|
||||
* @param dst[out] Destination tile where the result is stored.
|
||||
* @param row_values[in] Row vector containing values to broadcast into cols.
|
||||
*/
|
||||
template<typename ST, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
broadcast_col(threadgroup ST &dst, threadgroup const SV &col_values, const ushort laneid) {
|
||||
col_map<base_ops::copy2, ST, SV>(dst, dst, col_values, laneid);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,295 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Warp-scope reductions on shared tiles.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
|
||||
/**
|
||||
* Performs row-wise reduction on a matrix using a specified operation.
|
||||
*
|
||||
* @tparam op The operation to be applied for reduction.
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type with row layout.
|
||||
* @param row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param src The source matrix on which to perform the reduction.
|
||||
* @param src_accum The initial value of the accumulator, used when reset is false.
|
||||
* @param reset A boolean flag indicating whether to reset the accumulator (ignore src_accum) or not.
|
||||
*/
|
||||
template<typename op, typename SV, typename ST, bool reset>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_reduce(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
|
||||
using dtype = typename SV::dtype;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int row = laneid; row < ST::rows; row += mittens::SIMD_THREADS) {
|
||||
dtype accum = src[{row, 0}];
|
||||
#pragma clang loop unroll(full)
|
||||
for (int col = 1; col < src.cols; col++) {
|
||||
accum = op::template op<dtype>(accum, src[{row, col}]);
|
||||
}
|
||||
if (reset) {
|
||||
row_accum[row] = accum;
|
||||
} else {
|
||||
row_accum[row] = op::template op<dtype>(src_accum[row], accum);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs column-wise reduction on a matrix using a specified operation.
|
||||
*
|
||||
* @tparam op The operation to be applied for reduction.
|
||||
* @tparam V The shared vector type for the column accumulator.
|
||||
* @tparam T The shared matrix type with column layout.
|
||||
* @param col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param src The source matrix on which to perform the reduction.
|
||||
* @param src_accum The initial value of the accumulator, used when reset is false.
|
||||
* @param reset A boolean flag indicating whether to reset the accumulator (ignore src_accum) or not.
|
||||
*/
|
||||
template<typename op, typename SV, typename ST, bool reset>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_reduce(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
|
||||
using dtype = typename SV::dtype;
|
||||
#pragma clang loop unroll(full)
|
||||
for (int col = laneid; col < src.cols; col += mittens::SIMD_THREADS) {
|
||||
dtype accum = src[int2(0, col)];
|
||||
#pragma clang loop unroll(full)
|
||||
for (int row = 1; row < src.rows; row++) {
|
||||
accum = op::template op<dtype>(accum, src[int2(row, col)]);
|
||||
}
|
||||
if (reset) {
|
||||
col_accum[col] = accum;
|
||||
} else {
|
||||
col_accum[col] = op::template op<dtype>(src_accum[col], accum);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
|
||||
|
||||
/**
|
||||
* @brief Store the maximum of each row of the src shared matrix in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_max(threadgroup SV &row_accum, threadgroup const ST &src, const ushort laneid) {
|
||||
row_reduce<base_ops::max, SV, ST, true>(row_accum, src, row_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the minimum of each row of the src shared matrix in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_min(threadgroup SV &row_accum, threadgroup const ST &src, const ushort laneid) {
|
||||
row_reduce<base_ops::min, SV, ST, true>(row_accum, src, row_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the sum of each row of the src shared matrix in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_sum(threadgroup SV &row_accum, threadgroup const ST &src, const ushort laneid) {
|
||||
row_reduce<base_ops::sum, SV, ST, true>(row_accum, src, row_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the product of each row of the src shared matrix in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_prod(threadgroup SV &row_accum, threadgroup const ST &src, const ushort laneid) {
|
||||
row_reduce<base_ops::mul, SV, ST, true>(row_accum, src, row_accum, laneid);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store the maximum of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_max(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
|
||||
row_reduce<base_ops::max, SV, ST, false>(row_accum, src, src_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the minimum of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_min(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
|
||||
row_reduce<base_ops::min, SV, ST, false>(row_accum, src, src_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the sum of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_sum(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
|
||||
row_reduce<base_ops::sum, SV, ST, false>(row_accum, src, src_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the product of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] row_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
row_prod(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
|
||||
row_reduce<base_ops::mul, SV, ST, false>(row_accum, src, src_accum, laneid);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store the maximum of each column of the src shared matrix in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_max(threadgroup SV &col_accum, threadgroup const ST &src, const ushort laneid) {
|
||||
col_reduce<base_ops::max, SV, ST, true>(col_accum, src, col_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the minimum of each column of the src shared matrix in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_min(threadgroup SV &col_accum, threadgroup const ST &src, const ushort laneid) {
|
||||
col_reduce<base_ops::min, SV, ST, true>(col_accum, src, col_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the sum of each column of the src shared matrix in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_sum(threadgroup SV &col_accum, threadgroup const ST &src, const ushort laneid) {
|
||||
col_reduce<base_ops::sum, SV, ST, true>(col_accum, src, col_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the product of each column of the src shared matrix in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_prod(threadgroup SV &col_accum, threadgroup const ST &src, const ushort laneid) {
|
||||
col_reduce<base_ops::mul, SV, ST, true>(col_accum, src, col_accum, laneid);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Store the maximum of each column of the src shared matrix, as well as the src_accum shared vector, in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_max(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
|
||||
col_reduce<base_ops::max, SV, ST, false>(col_accum, src, src_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the minimum of each column of the src shared matrix, as well as the src_accum shared vector, in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_min(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
|
||||
col_reduce<base_ops::min, SV, ST, false>(col_accum, src, src_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the sum of each column of the src shared tile, as well as the src_accum row vector, in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_sum(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
|
||||
col_reduce<base_ops::sum, SV, ST, false>(col_accum, src, src_accum, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Store the product of each column of the src shared tile, as well as the src_accum row vector, in the col_accum shared vector.
|
||||
*
|
||||
* @tparam V The shared vector type for the row accumulator.
|
||||
* @tparam T The shared matrix type.
|
||||
* @param[out] col_accum The accumulator where the result of the reduction is stored.
|
||||
* @param[in] src The source matrix on which to perform the reduction.
|
||||
* @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
|
||||
*/
|
||||
template<typename SV, typename ST>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
|
||||
col_prod(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const ushort laneid) {
|
||||
col_reduce<base_ops::mul, SV, ST, false>(col_accum, src, src_accum, laneid);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
#pragma once
|
||||
#include "conversions.metal"
|
||||
#include "maps.metal"
|
||||
#include "reductions.metal"
|
||||
@@ -0,0 +1,60 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Warp-scope conversions on shared vectors.
|
||||
*/
|
||||
|
||||
#pragma once // done!
|
||||
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
|
||||
|
||||
/**
|
||||
* @brief Copies data from one shared vector to another, converting data types if necessary.
|
||||
*
|
||||
* This function copies data from the source shared vector `src` to the destination shared vector `dst`.
|
||||
* If the data types of `src` and `dst` are the same, it performs a direct memory copy. Otherwise, it
|
||||
* converts each element from the source data type to the destination data type using the appropriate
|
||||
* converter before copying.
|
||||
*
|
||||
* @tparam SV1 The type of the destination shared vector, must satisfy the ducks::sv::all concept.
|
||||
* @tparam SV2 The type of the source shared vector, must satisfy the ducks::sv::all concept.
|
||||
* @param[out] dst The destination shared vector.
|
||||
* @param[in] src The source shared vector.
|
||||
* @note The lengths of `src` and `dst` must be equal. This is enforced at compile time.
|
||||
*/
|
||||
template<typename SV1, typename SV2>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV1>() && ducks::is_shared_vector<SV2>(), void>::type
|
||||
copy(threadgroup SV1 &dst, threadgroup const SV2 &src, const ushort laneid) {
|
||||
static_assert(SV1::length == SV2::length, "Source and destination vectors must have the same length.");
|
||||
#pragma clang loop unroll(full)
|
||||
for(int i = laneid; i < dst.length; i+=SIMD_THREADS) {
|
||||
dst[i] = base_types::convertor<typename SV1::dtype, typename SV2::dtype>::convert(src[i]);
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------- SUBVEC ---------- */
|
||||
|
||||
/**
|
||||
* @brief Returns a reference to a subvec of a given shared vector
|
||||
*
|
||||
* @tparam subvec_tiles The length, in subtiles, of the subvec.
|
||||
* @tparam SV The type of the input vector, which must satisfy the ducks::sv::all concept.
|
||||
* @param src The input tile.
|
||||
* @param vec_idx The index of the subtile, in units of subvec_tiles*16 elements.
|
||||
* @return A reference to the subvec.
|
||||
*
|
||||
* @note The subvec length must evenly divide the vector length.
|
||||
*/
|
||||
template<int subvec_tiles, typename SV>
|
||||
//using subvec = typename SV::template subvec<SV::length>;
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), threadgroup typename SV::template subvec<typename SV::dtype, subvec_tiles>&>::type
|
||||
subvec_inplace(threadgroup SV &src, int vec_idx) {
|
||||
return *(threadgroup typename SV::template subvec<typename SV::dtype, subvec_tiles>*)(&src[vec_idx*TILE_DIM*subvec_tiles]);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,278 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Warp-scope maps on shared vectors.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
|
||||
/**
|
||||
* @brief Applies a unary operation to each element of a shared memory vector.
|
||||
*
|
||||
* @tparam op Unary operation type.
|
||||
* @tparam T Shared memory vector type.
|
||||
* @param dst[out] Destination vector in which to store the result.
|
||||
* @param src[in] Source vector to apply the unary operation.
|
||||
*/
|
||||
template<typename op, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
unary_op(threadgroup SV &dst, threadgroup const SV &src, const ushort laneid) {
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
#pragma clang loop unroll(full)
|
||||
for(int cur = laneid; cur < SV::length; cur+=SIMD_THREADS) {
|
||||
dst[cur] = op::template op<typename SV::dtype>(src[cur]);
|
||||
}
|
||||
}
|
||||
/**
|
||||
* @brief Perform a binary operation on two shared vectors.
|
||||
*
|
||||
* @tparam op The binary operation to perform.
|
||||
* @tparam T The type of the vectors.
|
||||
* @param dst[out] The destination vector where the result is stored.
|
||||
* @param lhs[in] The left-hand side vector for the operation.
|
||||
* @param rhs[in] The right-hand side vector for the operation.
|
||||
*/
|
||||
template<typename op, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
bin_op(threadgroup SV &dst, threadgroup const SV &lhs, threadgroup const SV &rhs, const ushort laneid) {
|
||||
#pragma clang loop unroll(full)
|
||||
for(int cur = laneid; cur < SV::length; cur+=SIMD_THREADS) {
|
||||
dst[cur] = op::template op<typename SV::dtype>(lhs[cur], rhs[cur]);
|
||||
}
|
||||
}
|
||||
/**
|
||||
* @brief Perform a binary operation on a shared vector and a scalar.
|
||||
*
|
||||
* @tparam op The binary operation to perform.
|
||||
* @tparam T The type of the vector.
|
||||
* @param dst[out] The destination vector where the result is stored.
|
||||
* @param src[in] The source vector for the operation.
|
||||
* @param param[in] The scalar parameter for the operation.
|
||||
*/
|
||||
template<typename op, typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
bin_op(threadgroup SV &dst, threadgroup const SV &src, thread const typename SV::T ¶m, const ushort laneid) {
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
#pragma clang loop unroll(full)
|
||||
for(int cur = laneid; cur < SV::length; cur+=SIMD_THREADS) {
|
||||
dst[cur] = op::template op<typename SV::dtype>(src[cur], param);
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
|
||||
|
||||
// ---- const ops ----
|
||||
|
||||
/**
|
||||
* @brief Sets all elements of a shared memory vector to zero.
|
||||
*
|
||||
* @tparam T Shared memory vector type.
|
||||
* @param dst[out] Destination vector to be set to zero.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
zero(threadgroup SV &dst, const ushort laneid) {
|
||||
unary_op<base_ops::zero, SV>(dst, dst, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of a shared memory vector to one.
|
||||
*
|
||||
* @tparam T Shared memory vector type.
|
||||
* @param dst[out] Destination vector to be set to one.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
one(threadgroup SV &dst, const ushort laneid) {
|
||||
unary_op<base_ops::one, SV>(dst, dst, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of a shared memory vector to positive infinity.
|
||||
*
|
||||
* @tparam T Shared memory vector type.
|
||||
* @param dst[out] Destination vector to be set to positive infinity.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
pos_infty(threadgroup SV &dst, const ushort laneid) {
|
||||
unary_op<base_ops::pos_infty, SV>(dst, dst, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Sets all elements of a shared memory vector to negative infinity.
|
||||
*
|
||||
* @tparam T Shared memory vector type.
|
||||
* @param dst[out] Destination vector to be set to negative infinity.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
neg_infty(threadgroup SV &dst, const ushort laneid) {
|
||||
unary_op<base_ops::neg_infty, SV>(dst, dst, laneid);
|
||||
}
|
||||
|
||||
// ---- unary ops ----
|
||||
|
||||
/**
|
||||
* @brief Copies the elements from one shared vector to another.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the source vector.
|
||||
* @param dst[out] Destination vector where the elements will be copied to.
|
||||
* @param src[in] Source vector to copy the elements from.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
copy(threadgroup SV &dst, thread const U &src, const ushort laneid) {
|
||||
bin_op<base_ops::copy2, SV>(dst, dst, src, laneid); // the second arg is ignored here.
|
||||
}
|
||||
/**
|
||||
* @brief Applies the exponential function element-wise to a shared vector.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @param dst[out] Destination vector where the exponential values will be stored.
|
||||
* @param src[in] Source vector to apply the exponential function to.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
exp(threadgroup SV &dst, threadgroup const SV &src, const ushort laneid) {
|
||||
unary_op<base_ops::exp, SV>(dst, src, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the exponential function element-wise to a shared vector, in base 2.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @param dst[out] Destination vector where the exponential values will be stored.
|
||||
* @param src[in] Source vector to apply the exponential function to.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
exp2(threadgroup SV &dst, threadgroup const SV &src, const ushort laneid) {
|
||||
unary_op<base_ops::exp2, SV>(dst, src, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the natural logarithm function element-wise to a shared vector.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @param dst[out] Destination vector where the logarithm values will be stored.
|
||||
* @param src[in] Source vector to apply the logarithm function to.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
log(threadgroup SV &dst, threadgroup const SV &src, const ushort laneid) {
|
||||
unary_op<base_ops::log, SV>(dst, src, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the absolute value function element-wise to a shared vector.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @param dst[out] Destination vector where the absolute values will be stored.
|
||||
* @param src[in] Source vector to apply the absolute value function to.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
abs(threadgroup SV &dst, threadgroup const SV &src, const ushort laneid) {
|
||||
unary_op<base_ops::abs, SV>(dst, src, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Applies the rectified linear unit (ReLU) function element-wise to a shared vector.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @param dst[out] Destination vector where the ReLU values will be stored.
|
||||
* @param src[in] Source vector to apply the ReLU function to.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
relu(threadgroup SV &dst, threadgroup const SV &src, const ushort laneid) {
|
||||
unary_op<base_ops::relu, SV>(dst, src, laneid);
|
||||
}
|
||||
|
||||
// ---- binary ops ----
|
||||
|
||||
/**
|
||||
* @brief Computes the element-wise maximum of two shared vectors.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the maximum values will be stored.
|
||||
* @param lhs[in] First vector for the maximum operation.
|
||||
* @param rhs[in] Second vector for the maximum operation.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
max(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const ushort laneid) {
|
||||
bin_op<base_ops::max, SV>(dst, lhs, rhs, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise minimum of two shared vectors.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the minimum values will be stored.
|
||||
* @param lhs[in] First vector for the minimum operation.
|
||||
* @param rhs[in] Second vector for the minimum operation.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
min(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const ushort laneid) {
|
||||
bin_op<base_ops::min, SV>(dst, lhs, rhs, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise sum of two shared vectors.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the sum values will be stored.
|
||||
* @param lhs[in] First vector for the sum operation.
|
||||
* @param rhs[in] Second vector for the sum operation.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
add(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const ushort laneid) {
|
||||
bin_op<base_ops::sum, SV>(dst, lhs, rhs, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise difference of two shared vectors.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the difference values will be stored.
|
||||
* @param lhs[in] First vector for the difference operation.
|
||||
* @param rhs[in] Second vector for the difference operation.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
sub(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const ushort laneid) {
|
||||
bin_op<base_ops::sub, SV>(dst, lhs, rhs, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise product of two shared vectors.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the product values will be stored.
|
||||
* @param lhs[in] First vector for the product operation.
|
||||
* @param rhs[in] Second vector for the product operation.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
mul(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const ushort laneid) {
|
||||
bin_op<base_ops::mul, SV>(dst, lhs, rhs, laneid);
|
||||
}
|
||||
/**
|
||||
* @brief Computes the element-wise division of two shared vectors.
|
||||
*
|
||||
* @tparam T Shared vector type.
|
||||
* @tparam U Type of the second vector.
|
||||
* @param dst[out] Destination vector where the division values will be stored.
|
||||
* @param lhs[in] First vector for the division operation.
|
||||
* @param rhs[in] Second vector for the division operation.
|
||||
*/
|
||||
template<typename SV, typename U>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
div(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const ushort laneid) {
|
||||
bin_op<base_ops::div, SV>(dst, lhs, rhs, laneid);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,268 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Warp-scope maps on shared vectors.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../../../../common/common.metal"
|
||||
#include "../../../../types/types.metal"
|
||||
|
||||
namespace mittens {
|
||||
|
||||
/**
|
||||
* @brief Performs a reduction operation on elements of a shared memory vector within a warp.
|
||||
*
|
||||
* This function applies a specified operation to reduce the elements of a shared memory vector `src` to a single value.
|
||||
* The result is stored in `accum`. If the `reset` parameter is true, the reduction includes an initial value `src_accum`.
|
||||
* The reduction operation is performed in a warp-wide context, ensuring synchronization between threads in the warp.
|
||||
*
|
||||
* @tparam op The operation to perform on the elements. Must provide a static `op` method.
|
||||
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
|
||||
* @tparam reset A boolean flag indicating whether to include an initial value in the reduction.
|
||||
* @param[out] accum The result of the reduction operation.
|
||||
* @param[in] src The shared memory vector to reduce.
|
||||
* @param[in] src_accum The initial value to include in the reduction if `reset` is false.
|
||||
*/
|
||||
template<typename op, typename SV, bool reset>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
reduce(thread typename SV::dtype &dst_accum, threadgroup const SV &src, thread const typename SV::dtype &src_accum, const ushort laneid) {
|
||||
using T = typename SV::dtype;
|
||||
|
||||
{
|
||||
T accum = src[0];
|
||||
for (int i = 1; i < SV::length; i++) {
|
||||
accum = op::template op<T>(accum, src[i]);
|
||||
}
|
||||
dst_accum = shfl_sync(accum, 0);
|
||||
return;
|
||||
}
|
||||
|
||||
//
|
||||
T accum;
|
||||
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
|
||||
for(int i = laneid + 32; i < SV::length; i+=32) {
|
||||
accum = op::template op<T>(accum, src[i]);
|
||||
}
|
||||
if (src.length >= 32) {
|
||||
// accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 1));
|
||||
accum = op::template op<T>(accum, (T)metal::simd_shuffle_rotate_down((float)accum, 1));
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
// accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 2));
|
||||
accum = op::template op<T>(accum, (T)metal::simd_shuffle_rotate_down((float)accum, 2));
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
// accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 4));
|
||||
accum = op::template op<T>(accum, (T)metal::simd_shuffle_rotate_down((float)accum, 4));
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
// accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 8));
|
||||
accum = op::template op<T>(accum, (T)metal::simd_shuffle_rotate_down((float)accum, 8));
|
||||
metal::simdgroup_barrier(metal::mem_flags::mem_none);
|
||||
// accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 16));
|
||||
accum = op::template op<T>(accum, (T)metal::simd_shuffle_rotate_down((float)accum, 16));
|
||||
|
||||
} else if (src.length == 24) {
|
||||
T shfl_val = shfl_down_sync<T>(accum, 1);
|
||||
accum = op::template op<T>(accum, shfl_val);
|
||||
|
||||
shfl_val = shfl_down_sync<T>(accum, 2);
|
||||
accum = op::template op<T>(accum, shfl_val);
|
||||
|
||||
shfl_val = shfl_down_sync<T>(accum, 4);
|
||||
accum = op::template op<T>(accum, shfl_val);
|
||||
|
||||
shfl_val = shfl_down_sync<T>(accum, 8);
|
||||
if (laneid < 16) {
|
||||
accum = op::template op<T>(accum, shfl_val);
|
||||
}
|
||||
shfl_val = shfl_down_sync<T>(accum, 16);
|
||||
accum = op::template op<T>(accum, shfl_val);
|
||||
} else if (src.length == 16) {
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 1));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 2));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 4));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 8));
|
||||
} else if (src.length == 8) {
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 1));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 2));
|
||||
accum = op::template op<T>(accum, shfl_down_sync<T>(accum, 4));
|
||||
}
|
||||
if (!reset) accum = op::template op<T>(accum, src_accum);
|
||||
dst_accum = shfl_sync(accum, 0);
|
||||
}
|
||||
|
||||
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
|
||||
|
||||
/**
|
||||
* @brief Finds the maximum element in a shared memory vector.
|
||||
*
|
||||
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
|
||||
* @param[out] max_val The maximum value found in the vector.
|
||||
* @param[in] src The shared memory vector to find the maximum in.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
max(thread typename SV::dtype &max_val, threadgroup const SV &src, const ushort laneid) {
|
||||
// reduce<base_ops::max, SV, true>(max_val, src, max_val, laneid);
|
||||
using T = typename SV::dtype;
|
||||
T accum = base_types::constants<T>::neg_infty();
|
||||
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
|
||||
for(int i = laneid + 32; i < SV::length; i+=32) {
|
||||
accum = base_ops::max::template op<T>(accum, src[i]);
|
||||
}
|
||||
max_val = (T)metal::simd_max((float)accum);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Finds the minimum element in a shared memory vector.
|
||||
*
|
||||
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
|
||||
* @param[out] min_val The minimum value found in the vector.
|
||||
* @param[in] src The shared memory vector to find the minimum in.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
min(thread typename SV::dtype &min_val, threadgroup const SV &src, const ushort laneid) {
|
||||
// reduce<base_ops::min, SV, true>(min_val, src, min_val);
|
||||
|
||||
using T = typename SV::dtype;
|
||||
T accum = base_types::constants<T>::pos_infty();
|
||||
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
|
||||
for(int i = laneid + 32; i < SV::length; i+=32) {
|
||||
accum = base_ops::min::template op<T>(accum, src[i]);
|
||||
}
|
||||
min_val = (T)metal::simd_min((float)accum);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Calculates the sum of elements in a shared memory vector.
|
||||
*
|
||||
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
|
||||
* @param[out] sum_val The sum of the values in the vector.
|
||||
* @param[in] src The shared memory vector to sum.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
sum(thread typename SV::dtype &sum_val, threadgroup const SV &src, const ushort laneid) {
|
||||
// reduce<base_ops::sum, SV, true>(sum_val, src, sum_val, laneid);
|
||||
using T = typename SV::dtype;
|
||||
T accum = base_types::constants<T>::zero();
|
||||
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
|
||||
for(int i = laneid + 32; i < SV::length; i+=32) {
|
||||
accum = base_ops::min::template op<T>(accum, src[i]);
|
||||
}
|
||||
sum_val = (T)metal::simd_sum((float)accum);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Calculates the product of elements in a shared memory vector.
|
||||
*
|
||||
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
|
||||
* @param[out] prod_val The product of the values in the vector.
|
||||
* @param[in] src The shared memory vector to multiply.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
prod(thread typename SV::dtype &prod_val, threadgroup const SV &src, const ushort laneid) {
|
||||
// reduce<base_ops::mul, SV, true>(prod_val, src, prod_val, laneid);
|
||||
using T = typename SV::dtype;
|
||||
T accum = base_types::constants<T>::one();
|
||||
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
|
||||
for(int i = laneid + 32; i < SV::length; i+=32) {
|
||||
accum = base_ops::min::template op<T>(accum, src[i]);
|
||||
}
|
||||
prod_val = (T)metal::simd_product((float)accum);
|
||||
}
|
||||
|
||||
// Three operand versions.
|
||||
|
||||
/**
|
||||
* @brief Finds the maximum element in a shared memory vector and accumulates it with src_accum.
|
||||
*
|
||||
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
|
||||
* @param[out] max_val The maximum value found in the vector, accumulated with src_accum.
|
||||
* @param[in] src The shared memory vector to find the maximum in.
|
||||
* @param[in] src_accum The initial value to accumulate with the maximum value found.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
max(thread typename SV::dtype &max_val, threadgroup const SV &src, thread const typename SV::dtype &src_accum, const ushort laneid) {
|
||||
// reduce<base_ops::max, SV, false>(max_val, src, src_accum, laneid);
|
||||
using T = typename SV::dtype;
|
||||
T accum = base_types::constants<T>::neg_infty();
|
||||
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
|
||||
for(int i = laneid + 32; i < SV::length; i+=32) {
|
||||
accum = base_ops::max::template op<T>(accum, src[i]);
|
||||
}
|
||||
max_val = (T)metal::simd_max((float)accum);
|
||||
max_val = base_ops::max::template op<T>(max_val, src_accum);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Finds the minimum element in a shared memory vector and accumulates it with src_accum.
|
||||
*
|
||||
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
|
||||
* @param[out] min_val The minimum value found in the vector, accumulated with src_accum.
|
||||
* @param[in] src The shared memory vector to find the minimum in.
|
||||
* @param[in] src_accum The initial value to accumulate with the minimum value found.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
min(thread typename SV::dtype &min_val, threadgroup const SV &src, thread const typename SV::dtype &src_accum, const ushort laneid) {
|
||||
// reduce<base_ops::min, SV, false>(min_val, src, src_accum, laneid);
|
||||
using T = typename SV::dtype;
|
||||
T accum = base_types::constants<T>::pos_infty();
|
||||
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
|
||||
for(int i = laneid + 32; i < SV::length; i+=32) {
|
||||
accum = base_ops::max::template op<T>(accum, src[i]);
|
||||
}
|
||||
min_val = (T)metal::simd_min((float)accum);
|
||||
min_val = base_ops::max::template op<T>(min_val, src_accum);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Calculates the sum of elements in a shared memory vector and accumulates it with src_accum.
|
||||
*
|
||||
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
|
||||
* @param[out] sum_val The sum of the values in the vector, accumulated with src_accum.
|
||||
* @param[in] src The shared memory vector to sum.
|
||||
* @param[in] src_accum The initial value to accumulate with the sum of the vector.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
sum(thread typename SV::dtype &sum_val, threadgroup const SV &src, threadgroup const typename SV::dtype &src_accum, const ushort laneid) {
|
||||
// reduce<base_ops::sum, SV, false>(sum_val, src, src_accum, laneid);
|
||||
using T = typename SV::dtype;
|
||||
T accum = base_types::constants<T>::zero();
|
||||
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
|
||||
for(int i = laneid + 32; i < SV::length; i+=32) {
|
||||
accum = base_ops::max::template op<T>(accum, src[i]);
|
||||
}
|
||||
sum_val = (T)metal::simd_sum((float)accum);
|
||||
sum_val = base_ops::max::template op<T>(sum_val, src_accum);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Calculates the product of elements in a shared memory vector and accumulates it with src_accum.
|
||||
*
|
||||
* @tparam SV The type of the shared memory vector. Must satisfy the `ducks::sv::all` concept.
|
||||
* @param[out] prod_val The product of the values in the vector, accumulated with src_accum.
|
||||
* @param[in] src The shared memory vector to multiply.
|
||||
* @param[in] src_accum The initial value to accumulate with the product of the vector.
|
||||
*/
|
||||
template<typename SV>
|
||||
static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
|
||||
prod(thread typename SV::dtype &prod_val, threadgroup const SV &src, thread const typename SV::dtype &src_accum, const ushort laneid) {
|
||||
// reduce<base_ops::mul, SV, false>(prod_val, src, src_accum, laneid);
|
||||
using T = typename SV::dtype;
|
||||
T accum = base_types::constants<T>::one();
|
||||
if(laneid < SV::length) accum = src[laneid]; // initialize a register accumulator
|
||||
for(int i = laneid + 32; i < SV::length; i+=32) {
|
||||
accum = base_ops::max::template op<T>(accum, src[i]);
|
||||
}
|
||||
prod_val = (T)metal::simd_product((float)accum);
|
||||
prod_val = base_ops::max::template op<T>(prod_val, src_accum);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
#pragma once
|
||||
#include "conversions.metal"
|
||||
#include "maps.metal"
|
||||
#include "reductions.metal"
|
||||
@@ -0,0 +1,4 @@
|
||||
#pragma once
|
||||
#include "memory/memory.metal"
|
||||
#include "register/register.metal"
|
||||
#include "shared/shared.metal"
|
||||
@@ -0,0 +1,4 @@
|
||||
#pragma once
|
||||
#include "common/common.metal"
|
||||
#include "ops/ops.metal"
|
||||
#include "types/types.metal"
|
||||
@@ -0,0 +1,63 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Templated layouts for complex global memory.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../../common/common.metal"
|
||||
//#include "../shared/cst.metal"
|
||||
#include "gl.metal"
|
||||
#include "util.metal"
|
||||
#ifdef mittens_HOPPER
|
||||
#include "tma.metal"
|
||||
#endif
|
||||
|
||||
namespace mittens {
|
||||
/* ---------- Global layout descriptor ---------- */
|
||||
|
||||
namespace ducks {
|
||||
namespace cgl {
|
||||
struct identifier {};
|
||||
}
|
||||
}
|
||||
|
||||
template<typename GL>
|
||||
struct cgl {
|
||||
static_assert(ducks::is_global_layout<GL>, "GL must satisfy global layout requirements.");
|
||||
|
||||
using identifier = ducks::cgl::identifier;
|
||||
using T = typename GL::T;
|
||||
using T2 = typename GL::T2;
|
||||
using dtype = typename GL::dtype;
|
||||
|
||||
GL real, imag;
|
||||
};
|
||||
|
||||
namespace ducks {
|
||||
template <typename T>
|
||||
struct has_cgl_identifier {
|
||||
static constant constexpr bool value = false; // Default case
|
||||
};
|
||||
|
||||
//template <typename _T, int b, int d, int r, int c, typename... TMA_Types>
|
||||
//struct has_cgl_identifier<mittens::gl<_T, b, d, r, c, TMA_Types ...>> {
|
||||
// static constant constexpr bool value = true;
|
||||
//};
|
||||
template <typename _T, int b, int d, int r, int c>
|
||||
struct has_cgl_identifier<mittens::gl<_T, b, d, r, c>> {
|
||||
static constant constexpr bool value = true;
|
||||
};
|
||||
|
||||
template <typename GL>
|
||||
static constexpr bool is_complex_global_layout() {
|
||||
return has_rt_identifier<GL>::value;
|
||||
}
|
||||
template <typename GL>
|
||||
static constexpr void assert_cgl() {
|
||||
static_assert(is_complex_global_layout<GL>(), "T must be a cgl");
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Templated layouts for global memory.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../../common/common.metal"
|
||||
#include "../shared/shared.metal"
|
||||
#include "../register/register.metal"
|
||||
#include "util.metal"
|
||||
|
||||
|
||||
namespace mittens {
|
||||
/* ---------- Associative dictionary for global layouts ---------- */
|
||||
|
||||
namespace detail {
|
||||
template<typename... Args>
|
||||
struct descriptor_dict {
|
||||
METAL_FUNC descriptor_dict() {}
|
||||
template<typename T> METAL_FUNC descriptor_dict(T _, int b, int d, int r, int c) {}
|
||||
METAL_FUNC descriptor_dict(thread const descriptor_dict &other) {}
|
||||
};
|
||||
}
|
||||
|
||||
/* ---------- Global layout descriptor ---------- */
|
||||
|
||||
namespace ducks {
|
||||
namespace gl {
|
||||
struct identifier {};
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static constexpr bool is_tile() {
|
||||
return mittens::ducks::is_shared_tile<T>() || mittens::ducks::is_register_tile<T>();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static constexpr bool is_vec() {
|
||||
return mittens::ducks::is_shared_vector<T>() || mittens::ducks::is_register_vector<T>();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template<typename _T, int b, int d, int r, int c>
|
||||
struct gl {
|
||||
using identifier = ducks::gl::identifier;
|
||||
|
||||
using T = typename base_types::packing<_T>::unpacked_type;
|
||||
using T2 = typename base_types::packing<_T>::packed_type;
|
||||
using dtype = T;
|
||||
|
||||
device T* raw_ptr;
|
||||
|
||||
ducks::g::make_dim_t<b> batch;
|
||||
ducks::g::make_dim_t<d> depth;
|
||||
ducks::g::make_dim_t<r> rows;
|
||||
ducks::g::make_dim_t<c> cols;
|
||||
// int batch;
|
||||
// int depth;
|
||||
// int rows;
|
||||
// int cols;
|
||||
|
||||
METAL_FUNC gl(device T *_data,
|
||||
ducks::g::make_arg_t<b> _batch,
|
||||
ducks::g::make_arg_t<d> _depth,
|
||||
ducks::g::make_arg_t<r> _rows,
|
||||
ducks::g::make_arg_t<c> _cols) :
|
||||
raw_ptr(_data), batch(_batch), depth(_depth), rows(_rows), cols(_cols) {
|
||||
}
|
||||
// METAL_FUNC gl(device T *_data,
|
||||
// int _batch,
|
||||
// int _depth,
|
||||
// int _rows,
|
||||
// int _cols) :
|
||||
// raw_ptr(_data), batch(_batch), depth(_depth), rows(_rows), cols(_cols) {
|
||||
// }
|
||||
//
|
||||
METAL_FUNC gl(thread const gl &other) :
|
||||
raw_ptr(other.raw_ptr), batch(other.batch), depth(other.depth), rows(other.rows), cols(other.cols) {}
|
||||
|
||||
METAL_FUNC gl(constant const gl &other) :
|
||||
raw_ptr(other.raw_ptr), batch(other.batch), depth(other.depth), rows(other.rows), cols(other.cols) {}
|
||||
|
||||
METAL_FUNC device T& operator[](const thread coord &idx) {
|
||||
return raw_ptr[((idx.b*depth + idx.d)*rows + idx.r)*cols + idx.c];
|
||||
}
|
||||
METAL_FUNC device const T& operator[](const thread coord &idx) const {
|
||||
return raw_ptr[((idx.b*depth + idx.d)*rows + idx.r)*cols + idx.c];
|
||||
}
|
||||
template<typename TILE>
|
||||
METAL_FUNC typename metal::enable_if<ducks::is_tile<TILE>(), device T&>::type
|
||||
get(const thread coord &idx) {
|
||||
return raw_ptr[((idx.b*depth + idx.d)*rows + idx.r*TILE::rows)*cols + idx.c*TILE::cols];
|
||||
}
|
||||
template<typename TILE>
|
||||
METAL_FUNC typename metal::enable_if<ducks::is_tile<TILE>(), device const T&>::type
|
||||
get(const thread coord &idx) const {
|
||||
return raw_ptr[((idx.b*depth + idx.d)*rows + idx.r*TILE::rows)*cols + idx.c*TILE::cols];
|
||||
}
|
||||
template<typename VEC>
|
||||
METAL_FUNC typename metal::enable_if<ducks::is_vec<VEC>(), device T&>::type
|
||||
get(const thread coord &idx) {
|
||||
return raw_ptr[((idx.b*depth + idx.d)*rows + idx.r)*cols + idx.c*VEC::length];
|
||||
}
|
||||
template<typename VEC>
|
||||
METAL_FUNC typename metal::enable_if<ducks::is_vec<VEC>(), device const T&>::type
|
||||
get(const thread coord &idx) const {
|
||||
return raw_ptr[((idx.b*depth + idx.d)*rows + idx.r)*cols + idx.c*VEC::length];
|
||||
}
|
||||
METAL_FUNC size_t row_stride() const { return cols; }
|
||||
};
|
||||
|
||||
namespace ducks {
|
||||
template <typename T>
|
||||
struct has_gl_identifier {
|
||||
static constant constexpr bool value = false; // Default case
|
||||
};
|
||||
|
||||
template <typename _T, int b, int d, int r, int c>
|
||||
struct has_gl_identifier<mittens::gl<_T, b, d, r, c>> {
|
||||
static constant constexpr bool value = true;
|
||||
};
|
||||
|
||||
template <typename GL>
|
||||
static constexpr bool is_global_layout() {
|
||||
return has_gl_identifier<GL>::value;
|
||||
}
|
||||
template <typename GL>
|
||||
static constexpr void assert_gl() {
|
||||
static_assert(is_global_layout<GL>(), "T must be a gl");
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
template<typename _T, int b, int d, int r, int c>
|
||||
struct gl2 {
|
||||
using identifier = ducks::gl::identifier;
|
||||
|
||||
using T = typename base_types::packing<_T>::unpacked_type;
|
||||
using T2 = typename base_types::packing<_T>::packed_type;
|
||||
using dtype = T;
|
||||
|
||||
device T* raw_ptr;
|
||||
|
||||
// ducks::g::make_dim_t<b> batch;
|
||||
// ducks::g::make_dim_t<d> depth;
|
||||
// ducks::g::make_dim_t<r> rows;
|
||||
// ducks::g::make_dim_t<c> cols;
|
||||
//
|
||||
// METAL_FUNC gl2(device T *_data,
|
||||
// ducks::g::make_arg_t<b> _batch,
|
||||
// ducks::g::make_arg_t<d> _depth,
|
||||
// ducks::g::make_arg_t<r> _rows,
|
||||
// ducks::g::make_arg_t<c> _cols) :
|
||||
// raw_ptr(_data), batch(_batch), depth(_depth), rows(_rows), cols(_cols) {
|
||||
// }
|
||||
|
||||
int batch;
|
||||
int depth;
|
||||
int rows;
|
||||
int cols;
|
||||
|
||||
METAL_FUNC gl2(device T *_data,
|
||||
int _batch,
|
||||
int _depth,
|
||||
int _rows,
|
||||
int _cols) :
|
||||
raw_ptr(_data), batch(_batch), depth(_depth), rows(_rows), cols(_cols) {
|
||||
}
|
||||
|
||||
|
||||
// METAL_FUNC gl2(thread const gl2 &other) :
|
||||
// raw_ptr(other.raw_ptr), batch(other.batch), depth(other.depth), rows(other.rows), cols(other.cols) {}
|
||||
//
|
||||
// METAL_FUNC gl2(constant const gl2 &other) :
|
||||
// raw_ptr(other.raw_ptr), batch(other.batch), depth(other.depth), rows(other.rows), cols(other.cols) {}
|
||||
|
||||
METAL_FUNC device T& operator[](const thread coord &idx) {
|
||||
return raw_ptr[((idx.b*depth + idx.d)*rows + idx.r)*cols + idx.c];
|
||||
}
|
||||
METAL_FUNC device const T& operator[](const thread coord &idx) const {
|
||||
return raw_ptr[((idx.b*depth + idx.d)*rows + idx.r)*cols + idx.c];
|
||||
}
|
||||
template<typename TILE>
|
||||
METAL_FUNC typename metal::enable_if<ducks::is_tile<TILE>(), device T&>::type
|
||||
get(const thread coord &idx) {
|
||||
return raw_ptr[((idx.b*depth + idx.d)*rows + idx.r*TILE::rows)*cols + idx.c*TILE::cols];
|
||||
}
|
||||
template<typename TILE>
|
||||
METAL_FUNC typename metal::enable_if<ducks::is_tile<TILE>(), device const T&>::type
|
||||
get(const thread coord &idx) const {
|
||||
return raw_ptr[((idx.b*depth + idx.d)*rows + idx.r*TILE::rows)*cols + idx.c*TILE::cols];
|
||||
}
|
||||
template<typename VEC>
|
||||
METAL_FUNC typename metal::enable_if<ducks::is_vec<VEC>(), device T&>::type
|
||||
get(const thread coord &idx) {
|
||||
return raw_ptr[((idx.b*depth + idx.d)*rows + idx.r)*cols + idx.c*VEC::length];
|
||||
}
|
||||
template<typename VEC>
|
||||
METAL_FUNC typename metal::enable_if<ducks::is_vec<VEC>(), device const T&>::type
|
||||
get(const thread coord &idx) const {
|
||||
return raw_ptr[((idx.b*depth + idx.d)*rows + idx.r)*cols + idx.c*VEC::length];
|
||||
}
|
||||
METAL_FUNC size_t row_stride() const { return cols; }
|
||||
};
|
||||
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief An aggregate header file for all the global types defined by Thundermittens.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
#include "util.metal"
|
||||
#include "gl.metal"
|
||||
#include "cgl.metal"
|
||||
@@ -0,0 +1,44 @@
|
||||
#pragma once
|
||||
|
||||
namespace mittens {
|
||||
namespace ducks {
|
||||
namespace g {
|
||||
|
||||
//template<int d> concept cdim = (d > 0); // represents a compile-time dimension
|
||||
//template<int d> concept rdim = (d == -1); // represents a runtime dimension
|
||||
|
||||
template<int d>
|
||||
struct compiled_dim {
|
||||
static_assert(d > 0, "Invalid compile-time dimension value"); // Replace `cdim` concept check
|
||||
static constant constexpr uint32_t v = d;
|
||||
|
||||
METAL_FUNC compiled_dim(thread const metal::nullptr_t &_) {}
|
||||
|
||||
METAL_FUNC constexpr operator uint32_t() const { return v; }
|
||||
};
|
||||
|
||||
struct runtime_dim {
|
||||
uint32_t v;
|
||||
METAL_FUNC runtime_dim(thread const uint32_t &_v) : v(_v) {}
|
||||
METAL_FUNC operator uint32_t() const { return v; }
|
||||
};
|
||||
|
||||
template<int d> using make_dim_t = metal::conditional_t<d == -1, runtime_dim, compiled_dim<d>>;
|
||||
template<int d> using make_arg_t = metal::conditional_t<d == -1, size_t, metal::nullptr_t>; // we pass runtime dims as size_t, comptime dims as nullptr_t
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
struct coord { // essentially a named int4 for tensor coordinates.
|
||||
int b, d, r, c;
|
||||
METAL_FUNC coord(int _b, int _d, int _r, int _c) : b(_b), d(_d), r(_r), c(_c) {}
|
||||
METAL_FUNC coord( int _d, int _r, int _c) : b( 0), d(_d), r(_r), c(_c) {}
|
||||
METAL_FUNC coord( int _r, int _c) : b( 0), d( 0), r(_r), c(_c) {}
|
||||
METAL_FUNC coord( int _c) : b( 0), d( 0), r( 0), c(_c) {}
|
||||
METAL_FUNC coord( ) : b( 0), d( 0), r( 0), c( 0) {}
|
||||
METAL_FUNC coord(thread const coord &other) : b(other.b), d(other.d), r(other.r), c(other.c) {}
|
||||
METAL_FUNC coord(thread const int4 &other) : b(other.x), d(other.y), r(other.z), c(other.w) {}
|
||||
METAL_FUNC operator int4() const { return int4(b, d, r, c); }
|
||||
};
|
||||
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Abstraction for a complex register tile composed of real and imaginary tiles
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "rt.metal"
|
||||
#include "crv.metal"
|
||||
|
||||
namespace mittens {
|
||||
|
||||
namespace ducks {
|
||||
namespace crt {
|
||||
/**
|
||||
* @brief A dummy type used to identify complex register tiles.
|
||||
*
|
||||
* For a type to quack like an rt_cmplx, it should define its identifier as ducks::rt::cmplx_identifier.
|
||||
* If a type quacks like ducks::rt::cmplx_identifier, it will be treated as an rt_cmplx by compiler checks.
|
||||
*/
|
||||
struct identifier {};
|
||||
} // namespace rt
|
||||
} // namespace ducks
|
||||
|
||||
/**
|
||||
* @brief Complex tile structure
|
||||
*
|
||||
* @tparam T2 The packed data type used for the matrix elements.
|
||||
* @tparam _rows The height of the tile in terms of the number of subtiles.
|
||||
* @tparam _cols The width of the tile in terms of the number of subtiles.
|
||||
* @tparam _layout The layout of the internal register tiles, either row-major or column-major.
|
||||
*
|
||||
* This structure is designed to abstract complex number operations internally to the real and imaginary
|
||||
* register tiles, respectively
|
||||
*
|
||||
* In general, you probably want a row-major tile, unless you specifically want to call mma
|
||||
*/
|
||||
template<typename _T, int _rows, int _cols, typename _layout>
|
||||
struct crt {
|
||||
using identifier = ducks::crt::identifier;
|
||||
static_assert(ducks::is_rt_layout<_layout>(), "crt was given invalid layout");
|
||||
using component = rt<_T, _rows, _cols, _layout>; /// Data type of each internal tile.
|
||||
using layout = typename component::layout; ///< Layout of the matrix tile, ensures compatibility with the rt concepts
|
||||
using T = typename component::T;
|
||||
using T2 = typename component::T2;
|
||||
using dtype = typename component::dtype; ///< Data type of the elements in the tile.
|
||||
|
||||
constant static constexpr int rows = component::rows;
|
||||
constant static constexpr int cols = component::cols;
|
||||
constant static constexpr int height = component::height;
|
||||
constant static constexpr int width = component::width;
|
||||
|
||||
// Real/imag tiles have same internal layout and size
|
||||
component real;
|
||||
component imag;
|
||||
|
||||
using row_vec = crv<T, cols, typename rt_base<T, layout>::row_vec_layout>; ///< A type representing a column vector for this tile.
|
||||
using col_vec = crv<T, rows, typename rt_base<T, layout>::col_vec_layout>; ///< A type representing a column vector for this tile.
|
||||
};
|
||||
|
||||
/* ---------- CONCEPTS ---------- */
|
||||
|
||||
namespace ducks {
|
||||
template <typename T>
|
||||
struct has_crt_identifier {
|
||||
static constant constexpr bool value = false; // Default case
|
||||
};
|
||||
|
||||
// Specialize for specific template instantiations of st
|
||||
template <typename _T, int _rows, int _cols, typename _layout>
|
||||
struct has_crt_identifier<mittens::crt<_T, _rows, _cols, _layout>> {
|
||||
static constant constexpr bool value = true;
|
||||
};
|
||||
|
||||
template <typename CRT>
|
||||
static constexpr bool is_complex_register_tile() {
|
||||
return has_crt_identifier<CRT>::value;
|
||||
}
|
||||
template <typename CRT>
|
||||
static constexpr void assert_complex_register_tile() {
|
||||
static_assert(is_register_tile<CRT>(), "T must be a rt");
|
||||
}
|
||||
}
|
||||
|
||||
template<int _rows, int _cols, typename _layout=ducks::rt_layout::row> using crt_fl = crt<float, _rows, _cols, _layout>;
|
||||
template<int _rows, int _cols, typename _layout=ducks::rt_layout::row> using crt_bf = crt<bf16, _rows, _cols, _layout>;
|
||||
template<int _rows, int _cols, typename _layout=ducks::rt_layout::row> using crt_hf = crt<half, _rows, _cols, _layout>;
|
||||
|
||||
|
||||
}
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Register vectors for computations on axes.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../../common/common.metal"
|
||||
#include "rv_layout.metal"
|
||||
#include "rv.metal"
|
||||
|
||||
namespace mittens {
|
||||
|
||||
/* ---------- MAIN VECTOR STRUCT ---------- */
|
||||
|
||||
// helper struct for type inference
|
||||
namespace ducks {
|
||||
/**
|
||||
* @namespace rt
|
||||
*
|
||||
* @brief The namespace where concepts and abstract types for register vectors live.
|
||||
*/
|
||||
namespace crv {
|
||||
/**
|
||||
* @brief A dummy type used to identify register vectors.
|
||||
*
|
||||
* For a type to quack like an rv, it should define its identifier as ducks::rv::identifier.
|
||||
* If a type quacks like ducks::rv::identifier, it will be treated as an rv by compiler checks.
|
||||
*/
|
||||
struct identifier {};
|
||||
}
|
||||
}
|
||||
/**
|
||||
* @brief Register vector structure.
|
||||
*
|
||||
* @tparam _T The packed data type used for the vector elements.
|
||||
* @tparam _outer_dim The size of the tile, in units of TILE_DIM (16).
|
||||
* @tparam _inner_dim This controls the layout of the tile in terms of which axis it maps on the register tile layout.
|
||||
*
|
||||
* Register vectors are used to accumulate and map values across tiles. You can do computation
|
||||
* on them directly if you want, but they're not designed to be maximally efficient vectors
|
||||
* as they have substantial duplication and strange layouts to help them work efficiently with
|
||||
* the register layouts used by the tensor cores. Thundermittens wants you working with tiles
|
||||
* where possible!
|
||||
*/
|
||||
|
||||
template<typename _T, size_t _length, typename _layout=ducks::rv_layout::naive>
|
||||
struct crv {
|
||||
static_assert(ducks::is_rv_layout<_layout>(), "_layout must be a rv layout");
|
||||
static_assert(ducks::base_types::isT1Type<_T>(), "T must be float, bf16, or half");
|
||||
using identifier = ducks::crv::identifier;
|
||||
using component = rv<_T, _length, _layout>; /// Data type of each internal tile.
|
||||
using layout = typename component::layout; ///< Layout of the matrix tile, ensures compatibility with the rv concepts
|
||||
|
||||
using T = typename component::T;
|
||||
using T2 = typename component::T2;
|
||||
using dtype = typename component::dtype; ///< Data type of the elements in the tile.
|
||||
|
||||
constant static constexpr int length = component::length;
|
||||
constant static constexpr int tiles = component::tiles;
|
||||
|
||||
// Real/imag tiles have same internal layout and size
|
||||
component real;
|
||||
component imag;
|
||||
};
|
||||
|
||||
/* ---------- CONCEPTS ---------- */
|
||||
|
||||
namespace ducks {
|
||||
template <typename T>
|
||||
struct has_crv_identifier {
|
||||
static constant constexpr bool value = false; // Default case
|
||||
};
|
||||
|
||||
// Specialize for specific template instantiations of st
|
||||
template <typename _T, int _length, typename _layout>
|
||||
struct has_crv_identifier<mittens::crv<_T, _length, _layout>> {
|
||||
static constant constexpr bool value = true;
|
||||
};
|
||||
|
||||
template <typename CRV>
|
||||
static constexpr bool is_complex_register_vector() {
|
||||
return has_crv_identifier<CRV>::value;
|
||||
}
|
||||
template <typename CRV>
|
||||
static constexpr void assert_complex_register_vector() {
|
||||
static_assert(is_complex_register_vector<CRV>(), "T must be a crv");
|
||||
}
|
||||
} // namespace ducks
|
||||
|
||||
template<int _l, typename layout=ducks::rv_layout::naive> using crv_fl = crv<float, _l, layout>;
|
||||
template<int _l, typename layout=ducks::rv_layout::naive> using crv_bf = crv<bf16, _l, layout>;
|
||||
template<int _l, typename layout=ducks::rv_layout::naive> using crv_hf = crv<half, _l, layout>;
|
||||
|
||||
|
||||
} // namespace mittens
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief An aggregate header file for all the register types defined by Thundermittens.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
#include "crv.metal"
|
||||
#include "rv.metal"
|
||||
#include "rv_layout.metal"
|
||||
#include "crt.metal"
|
||||
#include "rt.metal"
|
||||
#include "rt_layout.metal"
|
||||
#include "rt_base.metal"
|
||||
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief The main Thundermittens register tile struct, where most computation happens.
|
||||
*/
|
||||
#pragma once // kinda done
|
||||
/*
|
||||
TODO:
|
||||
consider if column layout rly rly rly makes no sense and no implement needed, not me being lazy
|
||||
*/
|
||||
#include <metal_stdlib>
|
||||
#include "../../common/common.metal"
|
||||
#include "rt_base.metal"
|
||||
#include "rv.metal"
|
||||
|
||||
/* ---------- MAIN TILE STRUCT ---------- */
|
||||
|
||||
|
||||
namespace mittens {
|
||||
/* ---------- MAIN TILE STRUCT ---------- */
|
||||
// helper struct for type inference
|
||||
namespace ducks {
|
||||
/**
|
||||
* @namespace rt
|
||||
*
|
||||
* @brief The namespace where concepts and abstract types for register tiles live.
|
||||
*/
|
||||
namespace rt {
|
||||
/**
|
||||
* @brief A dummy type used to identify register tiles.
|
||||
*
|
||||
* For a type to quack like an rt, it should define its identifier as ducks::rt::identifier.
|
||||
* If a type quacks like ducks::rt::identifier, it will be treated as an rt by compiler checks.
|
||||
*/
|
||||
struct identifier {};
|
||||
|
||||
} // namespace rt
|
||||
|
||||
} // namespace ducks
|
||||
|
||||
/**
|
||||
* @brief Main tile structure for manipulating data in registers.
|
||||
*
|
||||
* @tparam _T The data type used for the matrix elements.
|
||||
* @tparam _height The height of the tile in terms of the number of subtiles.
|
||||
* @tparam _width The width of the tile in terms of the number of subtiles.
|
||||
*
|
||||
* This structure is designed to handle matrix tiles in a flexible manner, allowing
|
||||
* for operations on tiles that are composed of smaller subtiles.
|
||||
*/
|
||||
template<typename _T, int _rows, int _cols, typename _layout=ducks::rt_layout::row>
|
||||
struct rt {
|
||||
using identifier = ducks::rt::identifier; ///< Type identifier for the rt structure.
|
||||
using layout = _layout;
|
||||
using T = typename base_types::packing<_T>::unpacked_type;
|
||||
static_assert(ducks::base_types::isT1Type<T>(), "T must be float, bf16, or half");
|
||||
static_assert(ducks::is_rt_layout<_layout>(), "T must be float, bf16, or half");
|
||||
using T2 = typename base_types::packing<_T>::packed_type;
|
||||
using dtype = T; ///< Data type of the elements in the tile.
|
||||
constant static constexpr int rows = _rows; ///< Total number of rows.
|
||||
static_assert(rows % rt_base<T, _layout>::tile_size == 0, "Rows must be divisible by the tile size");
|
||||
constant static constexpr int cols = _cols; ///< Total number of columns.
|
||||
static_assert(cols % rt_base<T, _layout>::tile_size == 0, "Columns must be divisible by the tile size");
|
||||
constant static constexpr int height = rows / rt_base<T, _layout>::tile_size; ///< Height in subtiles.
|
||||
constant static constexpr int width = cols / rt_base<T, _layout>::tile_size; ///< Width in subtiles.
|
||||
constant static constexpr int tile_size = rt_base<T, _layout>::tile_size; ///< Size of the base tile.
|
||||
constant static constexpr int num_elements = rt_base<T, _layout>::num_elements * width * height; ///< Total number of elements.
|
||||
constant static constexpr int elements_per_thread = rt_base<T, _layout>::elements_per_thread * width * height; ///< Elements handled per thread.
|
||||
constant static constexpr int packed_per_thread = rt_base<T, _layout>::packed_per_thread * width * height; ///< Packed elements per thread.
|
||||
constant static constexpr int packed_per_tile = rt_base<T, _layout>::packed_per_thread; ///< Packed elements per tile.
|
||||
|
||||
rt_base<dtype, _layout> tiles[height][width]; ///< The actual storage for the matrix tile, organized in subtiles.
|
||||
|
||||
using row_vec = rv<T, cols, typename rt_base<T, _layout>::row_vec_layout>; ///< A type representing a column vector for this tile.
|
||||
using col_vec = rv<T, rows, typename rt_base<T, _layout>::col_vec_layout>; ///< A type representing a column vector for this tile.
|
||||
};
|
||||
|
||||
|
||||
|
||||
namespace ducks{
|
||||
template <typename T>
|
||||
struct has_rt_identifier {
|
||||
static constant constexpr bool value = false; // Default case
|
||||
static constant constexpr bool is_row = false;
|
||||
static constant constexpr bool is_col = false;
|
||||
};
|
||||
|
||||
template <typename _T, int _rows, int _cols>
|
||||
struct has_rt_identifier<mittens::rt<_T, _rows, _cols, rt_layout::row>> {
|
||||
static constant constexpr bool value = true;
|
||||
static constant constexpr bool is_row = true; // Row-specific indicator
|
||||
static constant constexpr bool is_col = false;
|
||||
};
|
||||
|
||||
template <typename _T, int _rows, int _cols>
|
||||
struct has_rt_identifier<mittens::rt<_T, _rows, _cols, rt_layout::col>> {
|
||||
static constant constexpr bool value = true;
|
||||
static constant constexpr bool is_row = false;
|
||||
static constant constexpr bool is_col = true; // Col-specific indicator
|
||||
};
|
||||
|
||||
template <typename RT>
|
||||
static constexpr bool is_register_tile() {
|
||||
return has_rt_identifier<RT>::value;
|
||||
}
|
||||
|
||||
template <typename RT>
|
||||
static constexpr bool is_row_register_tile() {
|
||||
return has_rt_identifier<RT>::is_row;
|
||||
}
|
||||
|
||||
template <typename RT>
|
||||
static constexpr bool is_col_register_tile() {
|
||||
return has_rt_identifier<RT>::is_col;
|
||||
}
|
||||
|
||||
|
||||
template <typename RT>
|
||||
static constexpr void assert_register_tile() {
|
||||
static_assert(is_register_tile<RT>(), "T must be a rt");
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
|
||||
// layout and type wrappers
|
||||
|
||||
template<int _r, int _c, typename layout=ducks::rt_layout::row> using rt_fl = rt<float, _r, _c, layout>;
|
||||
template<int _r, int _c, typename layout=ducks::rt_layout::row> using rt_bf = rt<bf16, _r, _c, layout>;
|
||||
template<int _r, int _c, typename layout=ducks::rt_layout::row> using rt_hf = rt<half, _r, _c, layout>;
|
||||
} // namespace mittens
|
||||
@@ -0,0 +1,84 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief The basic 8x8 register tile on which larger register tiles are built.
|
||||
*/
|
||||
#pragma once // todo: col/row layout if needed
|
||||
#include <metal_stdlib>
|
||||
|
||||
#include "../../common/common.metal"
|
||||
#include "rt_layout.metal"
|
||||
#include "rv_layout.metal"
|
||||
namespace mittens {
|
||||
/* ---------- BASE 8x8 SUBTILE STRUCT ---------- */
|
||||
namespace ducks {
|
||||
/**
|
||||
* @namespace rt_base
|
||||
*
|
||||
* @brief The namespace where concepts and abstract types for register base (16x16) tiles live.
|
||||
*/
|
||||
namespace rt_base {
|
||||
/**
|
||||
* @brief A dummy type used to identify register base tiles.
|
||||
*
|
||||
* For a type to quack like an rt_base, it should define its identifier as ducks::rt_base::identifier.
|
||||
* If a type quacks like ducks::rt_base::identifier, it will be treated as an rt_base by compiler checks.
|
||||
*/
|
||||
struct identifier {};
|
||||
}
|
||||
template <typename T>
|
||||
static constexpr bool is_register_tile_base() {
|
||||
return metal::is_same<typename T::identifier, ducks::rt_base::identifier>::value;
|
||||
}
|
||||
template <typename RT>
|
||||
static constexpr void assert_register_tile_base() {
|
||||
static_assert(is_register_tile_base<RT>(), "T must be a rt_base");
|
||||
}
|
||||
} // namespace ducks
|
||||
|
||||
/**
|
||||
* @brief Basic tile structure for computation in registers.
|
||||
*
|
||||
* @tparam T2 The packed data type used for the matrix elements.
|
||||
* @tparam _layout The layout of the base tile, either row-major or column-major.
|
||||
*
|
||||
* This type is a primarily utility for building larger inline templates
|
||||
* out of PTX primitives and managing layouts.
|
||||
*
|
||||
* In general, you probably want a row-major tile, unless you specifically want to call mma
|
||||
*/
|
||||
template <typename _T, typename _layout>
|
||||
struct rt_base {
|
||||
using identifier = ducks::rt_base::identifier; ///< Type identifier for the rt_base structure.
|
||||
using layout = _layout; ///< Layout of the matrix tile.
|
||||
static_assert(ducks::base_types::isT1Type<_T>(), "rt_base was provided an unsupported type");
|
||||
static_assert(ducks::is_rt_layout<layout>(), "rt_base was provided an unsupported layout");
|
||||
using T = typename base_types::packing<_T>::unpacked_type;
|
||||
using T2 = typename base_types::packing<_T>::packed_type;
|
||||
using dtype = T;
|
||||
|
||||
|
||||
|
||||
static constant constexpr const int tile_size = mittens::TILE_DIM;
|
||||
static constant constexpr const int rows = tile_size;
|
||||
static constant constexpr const int cols = tile_size;
|
||||
static constant constexpr const int num_elements = rows*cols;
|
||||
static constant constexpr const int elements_per_thread = num_elements / mittens::SIMD_THREADS;
|
||||
|
||||
static constant constexpr const int registers_per_thread = elements_per_thread;
|
||||
static constant constexpr const int packed_per_thread = elements_per_thread / base_types::packing<T2>::num();
|
||||
metal::simdgroup_matrix<dtype, mittens::TILE_DIM, mittens::TILE_DIM> data;
|
||||
|
||||
using row_vec_layout = metal::conditional_t<metal::is_same_v<layout, ducks::rt_layout::row>, ducks::rv_layout::align, ducks::rv_layout::ortho>; // for holding column reductions
|
||||
|
||||
using col_vec_layout = metal::conditional_t<metal::is_same_v<layout, ducks::rt_layout::row>, ducks::rv_layout::ortho, ducks::rv_layout::align>; // for holding row reductions
|
||||
};
|
||||
|
||||
/* ---------- WRAPPERS FOR PRETTINESS ---------- */
|
||||
|
||||
template<typename L=ducks::rt_layout::row> using rt_base_fl = rt_base<float, L>;
|
||||
template<typename L=ducks::rt_layout::row> using rt_base_bf = rt_base<bf16, L>;
|
||||
template<typename L=ducks::rt_layout::row> using rt_base_hf = rt_base<half, L>;
|
||||
|
||||
|
||||
}
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Layouts and their manipulations for register tiles.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
|
||||
namespace mittens {
|
||||
namespace ducks {
|
||||
/**
|
||||
* @namespace rt_layout
|
||||
*
|
||||
* @brief A namespace for template metaprogramming with register tile layouts.
|
||||
*/
|
||||
namespace rt_layout {
|
||||
|
||||
/**
|
||||
* @brief A dummy type used to identify a row-major layout for a register tile.
|
||||
*/
|
||||
struct row {}; // for most matrices
|
||||
/**
|
||||
* @brief A dummy type used to identify a col-major layout for a register tile.
|
||||
*/
|
||||
struct col {}; // for the B-matrix of MMA ops.
|
||||
|
||||
template<typename l> struct transpose { using type = rt_layout::col; };
|
||||
template<> struct transpose<rt_layout::col> { using type = rt_layout::row; };
|
||||
} // namespace rt_layout
|
||||
template <typename _layout>
|
||||
METAL_FUNC static constexpr bool is_row_layout() {
|
||||
return metal::is_same_v<_layout, rt_layout::row>;
|
||||
}
|
||||
template <typename _layout>
|
||||
METAL_FUNC static constexpr bool is_col_layout() {
|
||||
return metal::is_same_v<_layout, rt_layout::col>;
|
||||
}
|
||||
template <typename _layout>
|
||||
METAL_FUNC static constexpr bool is_rt_layout() {
|
||||
return is_row_layout<_layout>() || is_col_layout<_layout>();
|
||||
}
|
||||
|
||||
|
||||
} // namespace ducks
|
||||
} // namespace mittens
|
||||
@@ -0,0 +1,125 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Register vectors for computations on axes.
|
||||
*/
|
||||
#pragma once
|
||||
#include "../../common/common.metal"
|
||||
#include "rv_layout.metal"
|
||||
namespace mittens {
|
||||
/* ---------- MAIN VECTOR STRUCT ---------- */
|
||||
|
||||
// helper struct for type inference
|
||||
namespace ducks {
|
||||
/**
|
||||
* @namespace rt
|
||||
*
|
||||
* @brief The namespace where concepts and abstract types for register vectors live.
|
||||
*/
|
||||
namespace rv {
|
||||
/**
|
||||
* @brief A dummy type used to identify register vectors.
|
||||
*
|
||||
* For a type to quack like an rv, it should define its identifier as ducks::rv::identifier.
|
||||
* If a type quacks like ducks::rv::identifier, it will be treated as an rv by compiler checks.
|
||||
*/
|
||||
struct identifier {};
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Register vector structure.
|
||||
*
|
||||
* @tparam _T The packed data type used for the vector elements.
|
||||
* @tparam _outer_dim The size of the tile, in units of TILE_DIM (8).
|
||||
* @tparam _inner_dim This controls the layout of the tile in terms of which axis it maps on the register tile layout.
|
||||
*
|
||||
* Register vectors are used to accumulate and map values across tiles. You can do computation
|
||||
* on them directly if you want, but they're not designed to be maximally efficient vectors
|
||||
* as they have substantial duplication and strange layouts to help them work efficiently with
|
||||
* the register layouts used by the tensor cores. Thundermittens wants you working with tiles
|
||||
* where possible!
|
||||
*/
|
||||
|
||||
template<typename _T, size_t _length, typename _layout>
|
||||
struct rv {
|
||||
using identifier = ducks::rv::identifier; ///< Type identifier for the rv structure.
|
||||
|
||||
static_assert(ducks::is_rv_layout<_layout>(), "_layout must be a rv layout");
|
||||
static_assert(ducks::base_types::isT1Type<_T>(), "T must be float, bf16, or half");
|
||||
using layout = _layout;
|
||||
constant static constexpr bool is_naive = ducks::is_naive_layout<layout>();
|
||||
using T = typename mittens::base_types::packing<_T>::unpacked_type;
|
||||
using T2 =typename mittens::base_types::packing<_T>::packed_type;
|
||||
using dtype = T; ///< Data type of the matrix elements
|
||||
|
||||
constant static constexpr int length = _length; ///< Length in elements.
|
||||
static_assert(length % mittens::TILE_DIM == 0, "Length must be divisible by the tile dimension");
|
||||
constant static constexpr int tiles = _length / mittens::TILE_DIM; ///< Length in subtiles, aliased for consistency with sv type
|
||||
constant static constexpr int inner_dim = layout::inner_dim; ///< Internal layout within a subtile. Either 1 or 2.
|
||||
constant static constexpr int outer_dim = is_naive ? (tiles+3)/4 : tiles; ///< Outer dim (also length in tiles)
|
||||
dtype data[outer_dim][inner_dim]; ///< The actual register vector data.
|
||||
|
||||
METAL_FUNC thread dtype* operator[](size_t idx) { return &data[idx][0]; } ///< A wrapper for indexing into vector data.
|
||||
METAL_FUNC thread const dtype* operator[](size_t idx) const { return &data[idx][0]; } ///< A wrapper for indexing into vector data.
|
||||
METAL_FUNC thread dtype& operator[](int2 outin) { return data[outin.x][outin.y]; } ///< A wrapper for indexing into vector data.
|
||||
METAL_FUNC thread const dtype& operator[](int2 outin) const { return data[outin.x][outin.y]; } ///< A wrapper for indexing into vector data.
|
||||
};
|
||||
|
||||
namespace ducks{
|
||||
template <typename T>
|
||||
struct has_rv_align_identifier {
|
||||
static constant constexpr bool value = false; // Default case
|
||||
};
|
||||
template <typename _T, int _length>
|
||||
struct has_rv_align_identifier<mittens::rv<_T, _length, ducks::rv_layout::align>> {
|
||||
static constant constexpr bool value = true;
|
||||
};
|
||||
template <typename RT>
|
||||
static constexpr bool is_align_register_vector() {
|
||||
return has_rv_align_identifier<RT>::value;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct has_rv_ortho_identifier {
|
||||
static constant constexpr bool value = false; // Default case
|
||||
};
|
||||
template <typename _T, int _length>
|
||||
struct has_rv_ortho_identifier<mittens::rv<_T, _length, ducks::rv_layout::ortho>> {
|
||||
static constant constexpr bool value = true;
|
||||
};
|
||||
|
||||
template <typename RT>
|
||||
static constexpr bool is_ortho_register_vector() {
|
||||
return has_rv_ortho_identifier<RT>::value;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct has_rv_naive_identifier {
|
||||
static constant constexpr bool value = false; // Default case
|
||||
};
|
||||
template <typename _T, int _length>
|
||||
struct has_rv_naive_identifier<mittens::rv<_T, _length, ducks::rv_layout::naive>> {
|
||||
static constant constexpr bool value = true;
|
||||
};
|
||||
template <typename RT>
|
||||
static constexpr bool is_naive_register_vector() {
|
||||
return has_rv_naive_identifier<RT>::value;
|
||||
}
|
||||
|
||||
template <typename RT>
|
||||
static constexpr bool is_register_vector() {
|
||||
return is_align_register_vector<RT>() || is_ortho_register_vector<RT>() || is_naive_register_vector<RT>();
|
||||
}
|
||||
|
||||
template <typename RT>
|
||||
static constexpr void assert_register_vector() {
|
||||
static_assert(is_register_vector<RT>(), "T must be a rv");
|
||||
}
|
||||
}
|
||||
template<int _l, typename layout=ducks::rv_layout::naive> using rv_fl = rv<float, _l, layout>;
|
||||
template<int _l, typename layout=ducks::rv_layout::naive> using rv_bf = rv<bf16, _l, layout>;
|
||||
template<int _l, typename layout=ducks::rv_layout::naive> using rv_hf = rv<half, _l, layout>;
|
||||
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user