Compare commits

...
Author SHA1 Message Date
geohot f082cbcb36 respect the 8x8 tiles 2025-10-07 17:51:24 +08:00
geohot 5ad62f130d split on tile_dim 2025-10-07 17:40:00 +08:00
geohot f129d75ee5 fix on cpu 2025-10-07 16:43:52 +08:00
geohot 51f3a5cbb4 gpu 2025-10-07 16:00:10 +08:00
geohot 1d7a8b33c1 gemm works on pc 2025-10-07 15:52:00 +08:00
George HotzandGitHub 3fae886aa9 Merge branch 'master' into clone_tg 2025-10-07 14:02:36 +08:00
George HotzandGitHub 514d2a0774 merge tagless reshapes (#12474)
* merge tagless reshapes

* cleanup
2025-10-07 13:57:58 +08:00
chenyuandGitHub 7b48f3cc45 failed test case repro for openpilot model (#12475)
* failed test case repro for openpilot model

* assertEqual
2025-10-07 13:46:43 +08:00
George HotzandGitHub 3f44ef699f Merge branch 'master' into clone_tg 2025-10-07 13:08:35 +08:00
geohot fa23f37e33 clone thunderkittens in uops 2025-10-07 13:06:58 +08:00
geohot 284db26a12 cleanup 2025-10-07 12:13:01 +08:00
geohot 0a0cb0b9e8 merge tagless reshapes 2025-10-07 12:11:11 +08:00
chenyuandGitHub a5484b767e remove skipping cast in simplify_valid [pr] (#12472)
* remove skipping cast in simplify_valid [pr]

unsupported statements are handled in uop_given_valid already. the test failed because (100%x) somehow got simplified

* better test
2025-10-07 00:10:04 -04:00
George HotzandGitHub b4509fba31 thundermittens (#12471)
* thundermittens

* give device a type
2025-10-07 11:47:39 +08:00
George HotzandGitHub 0f25b4b289 move frontend dir to nn [pr] (#12470) 2025-10-07 10:42:22 +08:00
qazalandGitHub f664bcc8bd use recursive_property in UOp tracing (#12469)
* test

* simple passing
2025-10-06 21:10:52 +03:00
qazalandGitHub 1af05dae77 fix rangeify in compile4.py (#12467)
* fix rangeify in compile4.py

* fix type_verify
2025-10-06 13:37:46 +03:00
qazalandGitHub 76e8a3250c rangeify: late zero folding (#12464)
* rangeify: late zero folding

* early

* not kernels

* none

* multi

* linter

* mstack is sink comment

* more comment
2025-10-06 12:52:33 +03:00
George HotzandGitHub 0c015a24fe use recursive_property to prevent RecursionError (#12465)
* use recursive_property to prevent RecursionError

* not slower

* fix tests

* faster

* simpler
2025-10-06 15:59:18 +08:00
chenyuandGitHub a1881b0c17 update test_chicken (#12466)
logits are close, just numerical
2025-10-06 03:58:44 -04:00
qazalandGitHub 1b1978b9c0 early copy fixup (#12463)
* simple failing test

* early copy fixup
2025-10-06 06:38:29 +03:00
chenyuandGitHub c1e85f699c multi test case for sharded ring allreduce (#12462)
* multi test case for sharded ring allreduce

triggers `children not making progress` with RANGEIFY

* expect_rangeify_fails
2025-10-05 23:18:24 -04:00
chenyuandGitHub 1823a5043f don't check MAX_BUFFER_SIZE on NULL (#12461) 2025-10-05 22:09:29 -04:00
George HotzandGitHub 46e8ea15c1 split pm_substitute_recurse (#12460) 2025-10-05 21:35:50 -04:00
nimlgenandGitHub 1216fff781 remote: raise runtimeerror in checkz (#12453) 2025-10-05 21:22:53 +08:00
qazalandGitHub 6ad9a688ed add failing test after "pend substitutes for speed" (#12457)
* add failing substitute test

* expect_rangeify_fails
2025-10-05 16:10:04 +03:00
chenyuandGitHub 74b04f7dca test beautiful_mnist_multigpu (#12455)
* test beautiful_mnist_multigpu

another example that fails with RANGEIFY

* now i remember

* MAX_BUFFER_SIZE=0
2025-10-05 08:45:01 -04:00
69857d0ab0 Stable Diffusion mlperf training (#11304)
* entrypoint for sd mlperf train development

* match sd-v2 mlperf reference unet

* implement dataloader from mlperf ref

* update dataloader reference

* implement LambdaLR scheduler from mlperf ref

* match tokenizer from mlperf reference

* sample latent

* add noise to latent

* complete training epoch

* run full training step

* jit training loop

* replicate mlperf ref. losses over 11 train steps

* save tinygrad loss checkpoints properly

* match out.2.bias.grad to reference

* match weights to ref after 1 step

* compare out.2.bias to ref over three train steps

* implement attn_mask; cleanup closeness testing

* correct mse loss

* update dev_run / dependencies

* setup validation config/checkpointing

* implement validation sampling

* test closeness of eval denoise step to mlperf ref

* test closeness of decoder to mlperf ref

* confirm inception matches mlperf ref

* resize w/ bicubic interpolation, test closeness

* confirm closeness of clip preprocess to mlperf ref

* confirm clip score matches mlperf ref

* confirm fid/clip scores match mlperf ref

* cleanup

* cleanup

* zero-init some unet params as in mlperf reference

* revert jit change

* uncomment dependencies

* move to tinybox red

* implement GradScaler from torch but jittable

* simplify lr_scheduler, ensure jittability

* instantiate GradScaler

* only check if grads are finite with fp16

* implement fp16 training loop

* refactor UNet: norm, gelu, mixed precision

* refactor clip_tokenizer to enable versioning

* make fp16 attention closer to torch

* remove comparisons to torch fp16 attention

* add globvars.py for reference

* confirm closeness of fp16 unet forward to mlperf

* test norm closeness to torch with precast

* remeasure e2e with master attention

* more detailed softmax upcast comparison to torch

* parameterize softmax upcast in attention and unet

* use fp32 weights with autocast to fp16

* cleanup

* add data/checkpoint download script

* debug kernel timeout on AMD

* fix finite grads check; start multigpu

* pass numpy arrays from dataloader

* include text encoder in jit train step

* use int32 for tokens instead of int64

* prevent multi bug in reshape within clip

* corealize more, del refs before

* add more logging and wandb

* use erf gelu in clip encoder

* minor changes to train step and logging

* save checkpoints for eval or resuming

* add eval-only logic to training script

* multigpu eval

* remove PARALLEL=0

* cleanup

* pad eval batches of size < EVAL_BS

* workaround silent multigpu bug in jit

* cleanup

* tokenize captions

* verify correctness of multigpu eval

* cleanup

* verify correctness of grads in train step

* verify correctness of training (20 steps)

* don't shard in the training jit

* training settings

* minor cleanup

* overfit train w/ eval on 6 samples

* offload to enable combined train and eval

* download to raid; use local rclone

* misc changes for mi300x / logging

* refactor eval for larger BS, verify correctness

* cleanup

* ckpt resuming and remove eval cats

* eval BEAM config on mi300x and red

* resume eval after crash

* confirm eval correctness (one iteration, 6 samples)

* verify eval correctness at full scale

* cleanup correctness testing

* training correctness (20 steps, BS=248 uniform)

* cleanup

* remove eval cache at end of run

* switch f16 for bf16, del grad scaler

* confirm bf16 training correctness

* timestamps, new jits

* merge jits in training

* realize loss/lr on CPU

* training correctness

* post-bf16 train/eval

* implement grad_acc with timing/logging

* beam offline; debug gradacc; use float32

* fix gradacc in jit, correctness test

* prepare f32 BS=512 gradacc=4 run

* workaround jit problem in diffusion eval

* scale lr by BS

* revert gradacc, prepare bf16 BS=336 lr*=BS train

* make checkpointing faster

* resume bf16 BS=336 base_lr=1.25e-7 run

* jit ckpt at beginning

* don't alloc more gpu mem in ckpt

* cleanup

* move script to mi300x dir

* cleanup

* cleanup unneeded files

* revert beam search to master

* minor changes

* fix regression: realize before assign in eval

* cleanup mlperf SD data/ckpt downloads

* workaround BEAM failure

* workaround bug in Tensor.stack

* minor changes

* revert gradscaler

* cleanup

* cleanup/validate dataloader

* ensure checksum of laion data

* simplify config

* load training state to jitted bufs

* simplify lr scheduler

* simplify train script

* cleanup comments

* refactor stable diffusion/unet init

* more refactoring of stable diffusion init

* fix import errors in tests

* refactor: separate train/eval

* fix import errors

* eval checkpoints in reverse chron. order

* save/load cycle in sd init

* refactor and verify eval

* verify training correctness

* prepare repro train run

* cleanup

* integrate beam retry, train, eval

* simplify wandb

* kill orphaned processes

* better logging

* train to 10 ckpts instead of 7

* remove optimizer/scheduler checkpointing/resume

* cleanup

* BEAM=2 7 ckpts

* add test to compare with torch softmax in amp

* cleanup

* stop eval early if checkpoint converged

* add test for lr scheduler

* add proper test method

* add test for training

* use venv name that is ignored by .gitignore

* linting

* add simple f32 softmax fxn

* revert change to scaled_dot_product_attention

* refactor gelu_erf init

* simplify mixed precision in unet

* add norm autocasting to fp32

* rm extra test

* test eval with NULL backend

* fix venv name

* simplify norm autocast

* use temp dir for training test

* actually add eval test

* remove parallel env variable from tests

* update clip with tests

* reorg init functions

* use np for testing

* remove unused var

* factor out GPUS

* add sd model init tests

* more unet tests

* match master

* rerun CI due to linux (remote) hang

* explain UNET_CKPTDIR

* rerun CI due to linux (remote) timeout

---------

Co-authored-by: chenyu <[email protected]>
2025-10-05 07:56:05 -04:00
George HotzandGitHub a976ace404 minor improvements to rewrite (#12454)
* minor improvements to rewrite

* need that continue

* faster
2025-10-05 18:09:32 +08:00
qazalandGitHub 4b60121498 fix bmnist torch with RANGEIFY=1 (#12442)
* fix bmnist torch with RANGEIFY=1

* alt

* test and comment

* this was always wrong

* simple failing test for rangeify

* simple upat to match the old behavior
2025-10-05 12:34:27 +03:00
George HotzandGitHub b5f31d7505 earlier seen children (#12451) 2025-10-05 15:55:13 +08:00
qazalandGitHub 865d5796f8 add a test for untested Tensor.assign behavior (#12448)
* add a test for untested Tensor.assign behavior

* better
2025-10-04 12:44:56 +03:00
Sieds LyklesandGitHub e74be4a140 UOp.factor and add chain sorting (#12413)
* add ordering

* fix some tests

* fix more tests

* shorten comment

* update test

* add rule and test

* add rule and test

* remove check

* use fold_divmod_congruence instead of simplify

* adjust tests

* shorten line

* new algo

* add test

* add function to un-nest the div

* add UOp.factor

* test UOp.factor

* uop_given_valid tries to factor simplex expression

* shorten line

* symbolic_flat is back

* change that back

* fix those new tests

* new rule for ordering

* factor multiple factors

* no symbolic_flat

* symbolic_flat to there

* move that back

* fix imports

* merge correctly

* linter happy

* add rule

* add a test

* cleanup

* revert that for now

* UOp.factor returns self instead of None

* try all_candidates

* remove or_else

* post index symbolic

* add test

* maket this closer to the original

* increase mac hlb_cifar min step time

* add some ordering tests

* cleanup

* increase pytest timeout time

* check dtype
2025-10-04 06:05:38 +02:00
Sieds LyklesandGitHub 394dc24110 post index symbolic (#12446)
* post index symbolic

* add test
2025-10-03 23:23:03 +02:00
chenyuandGitHub 9f2b69b870 enable few tests for PTX test_dtype (#12445) 2025-10-03 08:56:30 -04:00
George HotzandGitHub 0b534f71c2 recursive substitute should be O(n) (#12444)
* recursive substitute

* even faster

* make that a single rewrite
2025-10-03 18:29:59 +08:00
chenyuandGitHub b087663c35 RANGEIFY test_bert uses more ran somehow (#12443) 2025-10-03 04:38:53 -04:00
chenyuandGitHub 940a8d5ba9 default IGNORE_OOB=1 (#12441)
* default IGNORE_OOB=1

z3 can get very slow with RANGEIFY, also update some kernel numbers to what it is

* add to test
2025-10-03 04:16:19 -04:00
George HotzandGitHub d290e77a5b pend substitutes for speed (#12440) 2025-10-03 15:49:19 +08:00
nimlgenandGitHub 23d310bcc1 ptx: handle i8/u8 casts correctly (#12439)
* ptx: handle casts correctly

* notsetp
2025-10-03 15:34:15 +08:00
hoovedandGitHub 1e8945a28c Training loop for Stable Diffusion mlperf (#12315)
* add diff

* fix edit error

* match master

* point reference to specific commit

* simplify wandb logging

* remove lr test, dehardcode device

* increase stack size limit
2025-10-03 02:45:38 -04:00
George HotzandGitHub c7849ac593 fix test lil model (#12437)
* fix test lil model

* 4 not 3
2025-10-03 02:28:37 -04:00
chenyuandGitHub 0f82d92b9d use float for softmax in llm.py (#12438)
fixed numerical issue in `CPU=1 RANGEIFY=1 python3 -m tinygrad.apps.llm`
2025-10-03 02:27:56 -04:00
George HotzandGitHub 4c63f7e786 skip copies of reshaped buffers (#12430)
* skip copies of reshaped buffers

* always run NOOP

* comment

* comment
2025-10-03 13:05:47 +08:00
Sieds LyklesandGitHub 0047bcc535 undo loaded comparison swap (#12436)
* add rule

* add a test
2025-10-03 06:57:29 +02:00
chenyuandGitHub f203d8b221 update RANGEIFY kernel count and test_masked_select (#12435) 2025-10-03 00:41:34 -04:00
wozeparrotandGitHub a6dd5a224b skip webgpu tests (#12433) 2025-10-02 21:31:07 -07:00
chenyuandGitHub bf99de7b1e update a few more tests for RANGEIFY (#12434) 2025-10-03 00:16:58 -04:00
George HotzandGitHub 9cd365c12e little changes from double gemm (#12429)
* little changes from double gemm

* split pm_group_for_reduce

* pm_add_buffers_local

* Revert "pm_add_buffers_local"

This reverts commit 4d30a91db2.
2025-10-03 10:31:51 +08:00
Sieds LyklesandGitHub 16a65b4fd0 fix test_symbolic_gcd_div hang (#12427) 2025-10-03 04:21:16 +02:00
chenyuandGitHub 2d24af888b REWRITE_STACK_LIMIT (#12426) 2025-10-02 21:51:04 -04:00
hoovedandGitHub 1b58ef0d60 Increase stack size limit in unified_rewrite (#12424)
* increase stack size limit

* rerun CI due to random tqdm test fail
2025-10-03 09:06:47 +08:00
qazalandGitHub 17d36d0952 don't tag MSTACK/MSELECT on global buffers (#12423)
* don't tag MSTACK/MSELECT

* fix
2025-10-02 13:32:15 +03:00
chenyuandGitHub 7b3912d8e4 relax atol for some tests (#12422) 2025-10-02 05:04:44 -04:00
chenyuandGitHub 98163832e4 update RANGEIFY test_cast_padded (#12421)
* update RANGEIFY test_cast_padded

* update test
2025-10-02 04:37:35 -04:00
chenyuandGitHub 37beef6de3 add null bert training test in ci (#12420)
fails with RANGEIFY `RuntimeError: children not making progress`
2025-10-02 04:05:19 -04:00
f21851b099 ops: n^2 .device property fix (#12419)
* test case for a long rand chain

currently failing with RANGEIFY because device propogates too deep

* skip

* ops: n^2 .device property fix

* unskip

---------

Co-authored-by: Chen-Yu Yang <[email protected]>
2025-10-02 03:28:12 -04:00
b1tgandGitHub ec177c80c2 rangeify: fix test_where_fold (llvm) (#12416)
* rangeify: fix test_where_fold (AMD_LLVM)

* rm comment
2025-10-02 02:57:49 -04:00
qazalandGitHub 13a25b2e67 rangeify: don't shape INDEX on kernelize (#12417) 2025-10-02 09:45:37 +03:00
hoovedandGitHub 5d9035f5a6 Eval for Stable Diffusion mlperf (#12316)
* add diff

* rerun ci

* refactor beam workaround, add test

* fix conflict

* linting
2025-10-02 02:35:38 -04:00
0f804c9a83 Stable Diffusion model init for mlperf (#12314)
* include clip pr diff

* updated unet and sd init

* dehardcode default device

* revert beam hang workaround

---------

Co-authored-by: chenyu <[email protected]>
2025-10-02 02:28:41 -04:00
geohot 0eee93f0c0 hotfix: disable split ranges for non rangeify 2025-10-02 13:15:24 +08:00
George HotzandGitHub 583553f467 split ranges (#12411)
* split ranges

* simpler

* split ranges

* range str

* fix test

* oops

* faster

* no group 2

* tests

* dont_sub_ranges_for_image

* revert that
2025-10-02 12:57:22 +08:00
qazalandGitHub 6fc6b51b59 fix limit_bufs with kernelize (#12415) 2025-10-02 07:49:11 +03:00
qazalandGitHub d1c868f990 fix limit_bufs with multi (#12414) 2025-10-02 05:51:56 +03:00
qazalandGitHub 2fcd55583f allow less kernels in external_test_opt (#12412)
* allow less kernels in external_test_opt

* this was always 2
2025-10-02 05:05:42 +03:00
qazalandGitHub 8b48e19ce2 skip more multi remote tests (#12410) 2025-10-02 04:50:46 +03:00
geohot 3770dd9d80 annotate bufferize in viz 2025-10-02 09:20:50 +08:00
qazalandGitHub 5b649616ff rangeify: detect and assert cycles (#12405)
* rangeify: assert cycles

* rng=2

* any
2025-10-02 03:39:43 +03:00
Sieds LyklesandGitHub 9a64fc0d28 Load alt value with cast try 2 (#12407)
* add or_casted

* add tests and fix old tests

* cast load

* move that to pm_render

* add allow_any_len to gated load patterns in renderers

* slice [:2]
2025-10-02 00:55:29 +02:00
nimlgenandGitHub 3e0e0290ce increase timeout in test_module_runs (#12408) 2025-10-01 22:01:44 +03:00
Sieds LyklesandGitHub 2f8ac77c25 add allow_any_len to gated load patterns in renderers (#12406) 2025-10-01 20:35:32 +02:00
George HotzandGitHub 89bed28716 split reduceop (#12404)
* some rangeify tests fixed

* bring split reduceop to rangeify

* fix tests
2025-10-01 18:45:16 +08:00
George HotzandGitHub 74ee305948 some rangeify tests fixed (#12403) 2025-10-01 18:23:37 +08:00
qazalandGitHub f198a9e1ba skip test_multihost_aware_schedule, assign devices mismatch (#12396)
* minimal failing remote test

* this should've never worked?

* skip that test
2025-10-01 13:09:15 +03:00
ac3d457d5e rangeify: TestReduceOpsConstFolding (#12397)
Co-authored-by: George Hotz <[email protected]>
2025-10-01 17:58:19 +08:00
George HotzandGitHub 60e52fbe36 support opts in contig, simpler (#12400) 2025-10-01 17:20:04 +08:00
chenyuandGitHub 6c95b1f39d explicitly set device for CI unit test (#12399) 2025-10-01 05:16:54 -04:00
chenyuandGitHub 6ba8bf282f skip test_masked_select for RANGEIFY PYTHON (#12395) 2025-10-01 04:13:31 -04:00
chenyuandGitHub 689ab9151b more RANGEIFY tests (#12393)
would have caught the load alt regression without adding too many tests
2025-10-01 03:43:58 -04:00
chenyuandGitHub adc8c3b28f Revert "load alt value with cast (#12384)" (#12392)
This reverts commit 05e91a248d.
2025-10-01 03:20:04 -04:00
b1tgandGitHub 154d114364 rangeify: fix abstractions2.py (#12386)
* rangeify: fix abstractions2.py

* tests

* lint

* only abstractions2

* base
2025-10-01 09:58:56 +03:00
geohot fe96c8d345 add HALF flag to tinygrad.apps.llm 2025-10-01 14:44:59 +08:00
George HotzandGitHub f205352cd7 remove ranges with 1s (#12388)
* use op_in_parents

* remove the ranges of 1

* fix CL image thing

* fix realize
2025-10-01 14:43:29 +08:00
qazalandGitHub 90b1c0dd96 rangeify: test_where_fold kernel count (#12379)
* rangeify: test_where_fold kernel count

* get these from the index

* replace ranges

* fine

* movement ops

* diff

* better
2025-10-01 09:35:12 +03:00
b1tgandGitHub 42748ccb92 rangeify: fix test_prequant_conv2d_1x1 (#12391) 2025-10-01 02:33:47 -04:00
Sieds LyklesandGitHub 05e91a248d load alt value with cast (#12384)
* add or_casted

* add tests and fix old tests

* cast load

* move that to pm_render
2025-10-01 07:14:26 +02:00
qazalandGitHub 714500edfd viz: add font-weight to OffscreenCanvas config (#12390) 2025-10-01 08:08:47 +03:00
b1tgandGitHub 57ad46c6e4 rangeify: increase atol for test_two_binops_no_rerun passing on real windows machine (#12389)
CPU_LLVM=1
2025-10-01 00:56:45 -04:00
George HotzandGitHub e02da8f5ac use op_in_parents (#12385) 2025-10-01 12:37:29 +08:00
chenyuandGitHub 0662946fac atol in test_two_binops_no_rerun (#12387)
for RANGEIFY LLVM
2025-10-01 00:05:47 -04:00
b1tgandGitHub da52006bde rangeify: fix test_scatter_reduce (#12380)
* rangeify: fix test_scatter_reduce

* ext_vector_type

* set alignment=1 on boolean
2025-09-30 23:26:36 -04:00
George HotzandGitHub 1c1b4d14e9 minor cleaups in rangeify (#12382)
* minor cleaups in rangeify

* op_in_parents

* don't use toposort

* Revert "don't use toposort"

This reverts commit 257d8e2529.
2025-10-01 11:19:48 +08:00
wozeparrotandGitHub 4204edc60b feat: skip test_long (#12383) 2025-09-30 20:07:39 -07:00
chenyuandGitHub 8def8145e4 ALLOWED_KERNEL_COUNT openpilot 0.9.4 with RANGEIFY (#12381) 2025-09-30 22:58:59 -04:00
George HotzandGitHub 4c9a930de2 rangeify attn tests (#12377) 2025-10-01 09:59:19 +08:00
qazalandGitHub 26247573e1 rangeify multi tests on gpu (#12376)
* rangeify multi tests on gpu

* fix limit_bufs
2025-10-01 04:53:04 +03:00
qazalandGitHub f2eb92948d rangeify: ban view pushing (#12371)
* rangeify: ban view pushing

* don't shape INDEX

* fix the codegen cache

* make space
2025-10-01 04:37:52 +03:00
George HotzandGitHub a128fa0f8a removing double reshapes was wrong (#12375) 2025-10-01 09:25:35 +08:00
hoovedandGitHub 969a1b35ca LR scheduler for Stable Diffusion mlperf training (#12201)
* add lr scheduler for stable diffusion training

* add lr scheduler test

* rerun ci

* rerun CI

* use np for testing

* move test to CI path

* remove unneeded copy
2025-09-30 21:21:08 -04:00
George HotzandGitHub 9ef319f349 bad conv in rangeify (#12373)
* bad conv with broken rangeify

* no maxpool needed

* add empty_like

* typo

* no self

* issue remains for test
2025-10-01 08:56:22 +08:00
nimlgenandGitHub 080b26e7d7 use suppress_finalizing to not mute all exceptions (#12372) 2025-09-30 21:24:31 +03:00
George HotzandGitHub 44558a37f7 fix some rangeify tests (#12370)
* fix bad range merges

* fix rng

* fix uop gc

* fix some rangeify tests

* now that needs rangeify 2 also
2025-09-30 20:12:08 +08:00
nimlgenandGitHub 2c397eb2a2 rangeify: buf limit (#12336)
* limit bufs

* g

* fix buffer limit

* um?

* fix

* only these?

* typo

* f

* cleaner
2025-09-30 14:59:47 +03:00
George HotzandGitHub a83f219253 fix bad range merges (#12368)
* fix bad range merges

* fix rng

* fix uop gc
2025-09-30 19:30:21 +08:00
qazalandGitHub a95159d579 remove TestShapeSpec, it relies on ShapeTracker [pr] (#12369) 2025-09-30 14:20:35 +03:00
George HotzandGitHub 9cf5e66899 minimal rangeify stable diffusion fix (#12367)
* minimal rangeify stable diffusion fix

* more minimal
2025-09-30 18:48:35 +08:00
chenyuandGitHub b4a4817c9c fix rangeigy test_linalg (#12365) 2025-09-30 06:28:35 -04:00
qazalandGitHub de1d562b69 rangeify: update test_pickle asserts (#12366)
* realized exists on the base

* use is_realized
2025-09-30 13:27:41 +03:00
c9ef5d8fe5 rangeify: fix test_tensor_index_overflow (CPU_LLVM=1) (#12362)
* rangeify: fix test_tensor_index_overflow (CPU_LLVM=1)

* add test

---------

Co-authored-by: b1tg <[email protected]>
2025-09-30 05:55:15 -04:00
qazalandGitHub e8c595c29e remu: add new instructions introduced in RANGEIFY (#12363)
* add v_mad_i64_i32 for test_output_padded_conv_transpose2d

* run amd test_ops

* skip test_masked_select
2025-09-30 12:36:29 +03:00
George HotzandGitHub 360980f1a3 work on rangeify cost function heuristics (#12360)
* work on rangeify cost function heuristics

* dedup

* better cost function
2025-09-30 16:44:29 +08:00
qazalandGitHub 109c63b904 update Tensor unit tests for RANGEIFY (#12359)
* update test_kernelize for RANGEIFY

* also kernelizes user contiguous

* skip that test

* tensor uop repr

* 4 kernels, still realizes a float
2025-09-30 11:17:21 +03:00
George HotzandGitHub 7129419500 fix cifar training in RANGEIFY (#12355)
* fix cifar training in RANGEIFY

* even more wino fuse

* bugfix

* test to show issue
2025-09-30 15:59:19 +08:00
qazalandGitHub 4ff7f20b9d rangeify: fix kernelize (#12357) 2025-09-30 10:10:08 +03:00
chenyuandGitHub 86c5c969ea linalg cosmetic change (#12356) 2025-09-30 03:00:59 -04:00
qazalandGitHub 6a56d3c859 rangeify: only test correctness in multi (#12339)
* work

* more work

* back here

* skip tests

* work
2025-09-30 09:55:59 +03:00
George HotzandGitHub ab6b0d3a21 enable cleanup_dead_axes (#12351)
* enable cleanup_dead_axes

* don't mess with user contig

* correct tag behavior

* double reshape isn't correct

* block on assign too

* skip messing with symbolic

* Fix tests

* disable RANGEIFY=2

* test w rangeify
2025-09-30 14:09:39 +08:00
qazalandGitHub 2a7310ab59 rangeify: fix remaining multi correctness issue (#12354) 2025-09-30 08:08:27 +03:00
Sieds LyklesandGitHub 73b25bf47d z3 fix loaded mask (#12353)
* z3 fix loaded mask

* indentation
2025-09-30 06:55:50 +02:00
wozeparrotandGitHub 2a0caa09c2 push copy to disk (#12348) 2025-09-29 21:55:05 -07:00
chenyuandGitHub 881709cd33 don't skip rangeify test_instancenorm_3d (#12350)
seems fine now
2025-09-30 00:05:59 -04:00
hoovedandGitHub 39aae679e4 Support bfloat16 on NULL backend (#12340)
* add failing test

* move test

* only run test with NULL default

* add skip reason

* add fix
2025-09-30 00:02:30 -04:00
chenyuandGitHub af935e7d32 Revert "reduce const folding (#12344)" (#12349)
This reverts commit 8e508a9927.
2025-09-29 23:45:30 -04:00
George HotzandGitHub f522e83a02 fix rangeify elu fusion for openpilot (#12341)
* fix rangeify elu fusion for openpilot

* flip the metadata

* copy over permuted contiguous support

* this is correct

* update that
2025-09-30 11:41:52 +08:00
qazalandGitHub d95d018bb5 add name to multi rewrite [pr] (#12346) 2025-09-30 06:34:58 +03:00
qazalandGitHub 05275c9ec3 rangeify: enable assign to mstack target (#12345) 2025-09-30 06:27:57 +03:00
chenyuandGitHub 8e508a9927 reduce const folding (#12344) 2025-09-29 23:08:56 -04:00
chenyuandGitHub 3a480b858f use more getitem in gpt2 (#12343) 2025-09-29 23:08:03 -04:00
qazalandGitHub 32d69d07d7 rangeify: enable multitensor TestBatchNorm (#12342) 2025-09-30 06:05:00 +03:00
Sieds LyklesandGitHub d55d829635 Lower index dtype spec fix (#12337)
* new pm_lower_index_dtype

* load_store_indexing after index lowering

* shorten line

* seperate rule for long removal

* fix test

* fix index_to_concrete_int

* minor fixes

* add sink there

* update types in linearizer test
2025-09-30 04:26:50 +02:00
c38f6ce140 unified_rewrite: use deque and dont add nodes to the stack multiple times (#12320)
* use deque instead of list

* increase ctx.progress and max stack_len

* add openpilot

* prevent placing uops on stack many times

* revert increasing ctx.progress and stack length limit

* dont block adding to the stack there

---------

Co-authored-by: George Hotz <[email protected]>
2025-09-30 10:02:28 +08:00
hoovedandGitHub c2689c505e Clip model updates for Stable Diffusion mlperf training (#12313)
* stable diffusion mlperf clip changes

* add clip tests

* set gelu as attribute

* add more tests

* factor out GPUS

* rerun CI

* add imports to if blocks

* remove unneeded axis

* add clip tests to CI

* move clip tests

* add deps, disable max buf size
2025-09-29 21:50:14 -04:00
George HotzandGitHub cdfa0f29fd add rendering to index (#12338) 2025-09-30 09:18:05 +08:00
George HotzandGitHub baf3b60cfb fix gpt2 on rangeify (#12335) 2025-09-29 19:16:44 +08:00
qazalandGitHub 9513f025c5 apply multi before rangeify (#12298)
* it doesn't realize it when i reshape

* cleaner graph

* map out

* REDUCE_AXIS also gives the wrong answer

* maybe

* work

* back here

* try

* more

* refactor tests

* check MultiBuffer

* or copy

* fine with this

* don't need graph_rewrite_map in rangeify
2025-09-29 14:16:31 +03:00
George HotzandGitHub b899392f30 fix llm app with rangeify (#12334)
* fix llm app with rangeify

* add gpt2 contiguous also
2025-09-29 18:42:44 +08:00
wozeparrotandGitHub 7ae6898e31 better late bufferview (#12333) 2025-09-29 03:08:34 -07:00
George HotzandGitHub 3291e00df7 fix efficientnet slowness on rangeify (#12332) 2025-09-29 18:01:01 +08:00
chenyuandGitHub 9d2f2b8e34 skip test_mean_half_precision_overflow (#12331)
it only works with SPLIT_REDUCEOP=1
2025-09-29 05:15:04 -04:00
qazalandGitHub 9915bcf2b4 remove no-op contiguous from rand (#12329) 2025-09-29 11:53:16 +03:00
chenyuandGitHub 76c87d81b3 delete test_backward_sum_acc_dtype (#12330)
this test tests the wrong thing, it was only working because expand realize rule
2025-09-29 04:46:17 -04:00
George HotzandGitHub fd2e4f2353 failing rng test (#12328)
* tighten spec: fixup devectorizer types / rangeify

* tighten assign

* failing rangeify test

* simpler

* otherwise contig

* more tolerance cause rng seed changed
2025-09-29 16:06:45 +08:00
George HotzandGitHub 29469577e8 tighten spec: fixup devectorizer types / rangeify (#12327)
* tighten spec: fixup devectorizer types / rangeify

* tighten assign
2025-09-29 15:41:11 +08:00
wozeparrotandGitHub a982480512 feat: late to_bufferview (#12271) 2025-09-29 00:29:43 -07:00
qazalandGitHub e01a3eb59a rangeify whitespace cleanups [pr] (#12326)
* rangeify whitespace cleanups

* this is a noop
2025-09-29 10:04:51 +03:00
George HotzandGitHub cf925d1ac5 remove metadata for rangeify codegen (#12325) 2025-09-29 14:29:28 +08:00
George HotzandGitHub b252f890da add support for SPEC=1 (#12322)
* add support for SPEC=1

* cleaner place for it

* non rangeify spec

* split non rangeify
2025-09-29 12:55:01 +08:00
qazalandGitHub 292cb6ae26 viz: 404 if the requested rewrite doesn't exist (#12323) 2025-09-29 07:51:10 +03:00
qazalandGitHub 250cb10e8f rangeify permuted assign (#12299)
* enable RANGEIFY=1 test_assign

* work

* rangeify=0 asserts this ast

* remove that

* beta test, it's correct though

* skip multi

* matches torch/np output

* memcopy without memcopy

* can remove this

* rangeify isn't silently wrong anymore

* diff cleanup

* use UOp toposort instead of global tags

* actual assert TestRangeifyAssign

* step

* work

* this isn't optimizing away now

* some todos

* test fusion schedule

* typo

* dedup idxs

* cleaner

* pre

* work

* diff
2025-09-29 07:27:57 +03:00
Sieds LyklesandGitHub ed90de6583 Revert "Bufferize early, fix "children not making progress" on big graphs (#1…" (#12318)
This reverts commit 6f1cf717de.
2025-09-28 19:10:21 +02:00
Sieds LyklesandGitHub 29f0886395 skip test_softmax_fusion tests if RANGEIFY==1 (#12310) 2025-09-27 05:57:40 +02:00
Sieds LyklesandGitHub b98f1881ef dsp opt test has different axis number on rangeify (#12309) 2025-09-27 05:06:11 +02:00
Sieds LyklesandGitHub 6f1cf717de Bufferize early, fix "children not making progress" on big graphs (#12308)
* bufferize children early

* cleaner

* fix types

* lower number of reduceops

* test openpilot
2025-09-27 04:17:15 +02:00
qazalandGitHub 0104b16b9b rangeify: fix empty tags in reshapes (#12307) 2025-09-26 16:32:48 +03:00
nimlgenandGitHub f5eb46a3d9 fix limit buf metal on non rangeify (#12303)
* add failure test for limit buf on non rangeify

* correct metal

* correct

* hm
2025-09-26 11:06:28 +03:00
qazalandGitHub 8b2e0930d7 rangeify: enable passing multi test (#12301) 2025-09-26 08:31:13 +03:00
Sieds LyklesandGitHub 74411984fc Rangeify IMAGE (#12304)
* add imagedtype to rangeify

* enable some image tests

* move the tests

* image upcast before locals

* add if statement

* rangeify image_dtype test

* decrease read_image count
2025-09-26 07:21:02 +02:00
wozeparrotandGitHub d2cd269e28 fix: try close mmap (#12306) 2025-09-25 20:54:27 -07:00
chenyuandGitHub 17cec8d645 RANGEIFY winograd test (#12297)
speed seems fine
2025-09-24 23:42:32 -04:00
nimlgenandGitHub 476a2a0a96 test_qcom: update (#12293) 2025-09-24 21:45:58 +03:00
qazalandGitHub 38ecefaacb RANGEIFY=1 allreduce (#12260)
* ci

* extract mops

* work

* assert early

* port this?

* can realize shard

* allreduce passing

* notes

* better handling of shard

* err

* outerworld allreduce twice

* work

* don't tag movement ops

* don't tag movement ops

* delete old logic

* 19 failing + ram

* cleanup

* reset stuff

* simplest failing test

* diff

* test_ones

* allreduce work

* allreduce more work

* down to 22 failing tests

* port _device_num

* replace creates a new UOp here

* pour symbolic everywhere

* 7 failing

* focus on allreduce

* work

* cleanup

* more ci

* fix test_schedule_ring

* post index const shape

* much better

* diff cleanup
2025-09-24 18:13:08 +03:00
qazalandGitHub 0e778296be rangeify: refactor const folding (#12291)
* rangeify: refactor const folding [pr]

* it got better
2025-09-24 17:58:39 +03:00
qazalandGitHub 6c9d8c7e41 rangeify: simplify noop copy (#12289) 2025-09-24 17:01:23 +03:00
qazalandGitHub 1400ce105f rangeify: fix sharding (#12288) 2025-09-24 14:33:56 +03:00
qazalandGitHub 154c865966 rangeify: fix ram usage in multi (#12286) 2025-09-24 13:48:58 +03:00
Sieds LyklesandGitHub e8945c74de fix infinite symbolic loop with VCONST (#12285) 2025-09-24 07:06:22 +02:00
Sieds LyklesandGitHub 45c7252aed Better div nesting 2 (#11812)
* remove check

* use fold_divmod_congruence instead of simplify

* adjust tests

* shorten line

* new algo

* add test

* cleanup

* update tests

* ALLOWED_GATED_READ_IMAGE from 16 -> 12

* only remove the call to simplify

* add option to simplify with factor_remainder

* Allowed readimage gates back to 16
2025-09-24 04:50:26 +02:00
Sieds LyklesandGitHub 6146c64d81 lower the invalid gate last (#12164)
* lowering invalid gate is part of lower_index_dtype

* update test

* remove import

* put that back

* reduce_collapse uses invalid

* fix that pattern to use invalid_pat

* valid creates the right dtype count

* seperate rule for lowering invalid gate

* dont unvectorize Invalid gate

* image_fixup uses Invalid

* update tests

* cleanup

* update split_load_store

* add .scalar() there
2025-09-24 04:27:35 +02:00
qazalandGitHub ad7c8c21ea rangeify: INDEX doesn't passthrough MSELECT (#12279) 2025-09-23 21:36:50 +03:00
nimlgenandGitHub 02a7b7fe48 rangeify: fix test_setitem (#12269)
* rangeify: fix test_setitem

* um?

* better?

* simple where folding

* f

* revert

* x
2025-09-23 20:42:36 +03:00
qazalandGitHub 2f145a98e0 rangeify: fix contiguous multi (#12278)
* rangeify: fix contiguous multi

* when it's changing root, it should construct a new UOp
2025-09-23 20:05:29 +03:00
nimlgenandGitHub 5f4eeb054c rangeify: passes now (#12277) 2025-09-23 18:46:49 +03:00
qazalandGitHub 680ce54dd4 add types to replace_dnum (#12276) 2025-09-23 14:43:04 +03:00
chenyuandGitHub fffce0a6b4 use more no_range in simplify [pr] (#12275) 2025-09-23 02:33:56 -04:00
chenyuandGitHub 51b88b2265 process replay tests in rangeify (#12274) 2025-09-23 01:30:06 -04:00
chenyuandGitHub b54cb272d0 move test_qcom to test/device (#12272) 2025-09-22 21:07:10 -04:00
Sieds LyklesandGitHub d21e34e617 enable test_sum_twice (#12270)
* remove skip

* remove import
2025-09-23 00:57:29 +02:00
Sieds LyklesandGitHub 5a4b244e6b Check for group inside another reduce (#12268)
* add check

* get the ranges correctly

* add test

* comment and better check
2025-09-23 00:32:41 +02:00
qazalandGitHub a6fd96f620 rangeify: don't tag movement ops (#12267)
* don't tag movement ops

* delete old logic
2025-09-22 16:40:17 +03:00
chenyuandGitHub b03ceb806e move test_sample to test_randomness (#12266) 2025-09-21 21:11:32 -04:00
qazalandGitHub 25e0b725d1 cleanup section 0 rangeify (#12264) 2025-09-22 00:30:44 +03:00
qazalandGitHub 1aba668a37 cleanup buffer_view matcher (#12263) 2025-09-21 23:45:48 +03:00
nimlgenandGitHub b53a266254 rangeify: fix test_optim (#12262)
* rangeify: fix test_optim

* add to cl?

* these are good now
2025-09-21 18:08:35 +03:00
qazalandGitHub 461e9becec srender UOp in movement op arg (#12261) 2025-09-21 13:55:45 +03:00
Sieds LyklesandGitHub 9569fdfa36 use str for AxisType and AddrSpace __repr__ (#12252) 2025-09-21 05:24:41 +02:00
qazalandGitHub 8365c28cd5 viz: put a limit of brightness scale (#12259) 2025-09-20 18:52:55 +03:00
nimlgenandGitHub 4762a24022 test_free_intermediates force buffers (#12255)
* test_free_intermediates force buffers

* f

* fix for rangiefy

* xx
2025-09-20 18:14:39 +03:00
qazalandGitHub 57c7e0a8f8 RANGEIFY=1 test_jit (#12254)
* RANGEIFY=1 test_jit

* don't do any of that

* disk

* simple disk tensor

* more work

* run more tests

* it also doesn't copy everytime

* skip tests that hang everything
2025-09-20 17:34:32 +03:00
chenyuandGitHub 393c6b236c test case to sum twice in different order (#12253)
* test case to sum twice in different order

fixed by #12251

* try metal
2025-09-20 10:11:57 -04:00
qazalandGitHub 4756971c88 skip test_bf16_disk_write_read on CL=1 (#12256) 2025-09-20 17:11:06 +03:00
chenyuandGitHub 5e794be8af tighter spec for RANGE (#12250) 2025-09-20 07:59:50 -04:00
Sieds LyklesandGitHub 73c8dae60d add missing remove_blockend case (#12251)
* add missing remove_blockend case

* remove expectedFailure

* better comment
2025-09-20 06:29:19 +02:00
wozeparrotandGitHub dc4dd898b7 fix: close mmap (#12249) 2025-09-19 14:09:12 -07:00
Sieds LyklesandGitHub bb1f376ae6 profile z3 (#12248) 2025-09-19 22:52:06 +02:00
7e06d3ebba enable test_symbolic_jit (#12245)
Co-authored-by: qazal <[email protected]>
2025-09-19 20:23:42 +02:00
qazalandGitHub bb59eed82f rangeify: don't tag consts, they are global (#12247)
* rangeify: don't tag consts, they are global

* don't map movement ops

* sym failing test

* remove that

* update comment

* simpler test

* work
2025-09-19 15:25:03 +03:00
Sieds LyklesandGitHub cc038b31b6 Shrink instead of reshape to unregister symbolic (#12241)
* Slice to unbind symbolic

* use vmax for now

* assert shape in reshape is valid

* update test_symbolic_ops to use shrink instead of reshape

* remove infer_with_bound_values for npw

* symbolic output doesnt have symbolic strides

* symbolic jit tests use shrink to unregister symbolic

* update test

* update more tests

* wrap vmax in int()

* only create a new st if the store is not an assigne

* unwrap st

* comments
2025-09-19 06:04:35 +02:00
chenyuandGitHub a531a649fb test_resize_upsample_scales_cubic_align_corners_cpu is fixed (#12244) 2025-09-18 20:55:26 -04:00
Sieds LyklesandGitHub 8d703a6369 z3 xor doesnt use bitcast (#12243) 2025-09-19 00:31:44 +02:00
chenyuandGitHub 0dad6cc518 good RANGEIFY kernel counts in external_test_opt (#12242)
no push permute stuff. the model ones are less clear if it's good, some got slower
2025-09-18 17:58:54 -04:00
chenyuandGitHub cff1065f5e test CL=1 RANGEIFY=1 onnx (#12240)
all except test_resize_upsample_scales_cubic_align_corners_cpu runs
2025-09-18 16:49:46 -04:00
Sieds LyklesandGitHub ef05178855 fix 0//0 infinite rewrite in rangeify onnx (#12239) 2025-09-18 21:59:50 +02:00
chenyuandGitHub 87707ef0b8 unify range_start [pr] (#12236) 2025-09-18 13:52:54 -04:00
qazalandGitHub 825f148469 rangeify: fix copy size mismatch errs (#12232)
* rangeify: fix copy size mismatch errs

* const folding can happen in sym

assert it

* shippable

* rangeify copy is completely wrong

* pre_bufferize

* tag bufferize

* pre back
2025-09-18 18:23:32 +03:00
chenyuandGitHub f82b16a0e9 RANGEIFY test_tensor (#12235) 2025-09-18 10:35:43 -04:00
chenyuandGitHub 7487c13b61 truncate_fp16 -> float_to_fp16 (#12234)
match float_to_bf16 and float_to_fp8
2025-09-18 09:48:27 -04:00
54c15d74a4 python float8 support (#11960)
* basic support

* alu

* nan in exec_alu

* rand_for_dtype

* inf + 0.0

* finfo

* revert rand_for_dtype

* clean

* truncate fp8s inf

* spec ok

* float_to_fp8 nan/inf

* least_upper_dtype

* clean up

---------

Co-authored-by: b1tg <[email protected]>
2025-09-18 09:17:09 -04:00
qazalandGitHub dbbc261075 rangeify: fix COPY simplifier (#12233) 2025-09-18 14:35:33 +03:00
Sieds LyklesandGitHub f1108f1cbe Enable test_symbolic_ops on rangeify (#12230)
* enable

* merge correctly
2025-09-18 02:12:36 +02:00
Sieds LyklesandGitHub 812f485cd7 Enable threefry_doesnt_use_long test on rangeify (#12229)
* dont bufferize rangeify

* enable doesnt_use_long test
2025-09-18 01:58:34 +02:00
nimlgenandGitHub 3c5b8bf50c am: bump fw to rocm7 (#12226) 2025-09-17 21:20:22 +03:00
qazalandGitHub 525f80e0d2 rangeify: enable putting consts back in the tensor graph (#12225)
* rangeify: enable putting consts back in the tensor graph

* work

* sym in ci
2025-09-17 19:45:04 +03:00
chenyuandGitHub edffc246ed MUL in reduce_unparented (#12223)
* MUL in reduce_unparented

* some test
2025-09-17 11:56:39 -04:00
qazalandGitHub 7733c217c5 remove spam comments in test_schedule (#12224) 2025-09-17 18:24:55 +03:00
qazalandGitHub d917895569 map out rangeify errors in test_schedule (#12211)
* map out rangeify errors in test_schedule

* skip that

* add to ci
2025-09-17 09:10:28 +03:00
Sieds LyklesandGitHub 158506b91e Upgrade some divmod folding for symbolic divs (#12216)
* use const_factor() instead of arg

* add test

* change div min_max

* add tests

* add divide_by_symbolic_gcd

* add tests

* one more test

* Slice to unbind symbolic

* deal with const factor properly

* minor cleanup

* divide_by_symbolic_gcd becomes UOp.gcd and UOp.divide_exact

* add tests

* add gcd_without_const

* fix divide_exact bug

* add factor_remainder

* add tests

* fix imports

* elif -> if

* remove expectedFailure

* add more tests

* add more unwrap

* fix signature of pop_const

* remove that

* remove that
2025-09-17 03:00:50 +02:00
Sieds LyklesandGitHub 328bfe6b9b fix map_expand for symbolic shapes (#12218)
fix incorrect default argument in resolve
2025-09-17 01:20:18 +02:00
chenyuandGitHub 5b12764b83 add arange cat arange test (#12217)
simple test case to catch wrong reduce const folding. also clean up the old arange complexity test
2025-09-16 17:12:32 -04:00
nimlgenandGitHub 53655a4ee5 cuda: cleanup old comment (#12215) 2025-09-16 23:11:32 +03:00
chenyuandGitHub 6b808c5fe6 update TestSymbolicJit.test_plus1_pad (#12214)
was failing because movement was not captured
2025-09-16 15:57:50 -04:00
Shun UsamiandGitHub 2a72b00679 Add test for 2D tensor indexing in setitem (#12193)
* Add test for 2D tensor indexing in setitem

* Fix _masked_setitem to handle multi dim indexing correctly

* Fix indent

* Add fuzz test for 3D tensor indexing in setitem

* Skip indexing fuzz test (slow)
2025-09-16 14:57:25 -04:00
chenyuandGitHub c7b03457d7 Revert "Revert "more llvm intrinsics (#11961)" (#12194)" (#12195)
This reverts commit df1c183e46.
2025-09-16 14:55:31 -04:00
chenyuandGitHub 494bb12500 skip slow cifar bf16 on red benchmark (#12213)
very slow to compile the fake bf16
2025-09-16 14:55:01 -04:00
chenyuandGitHub 419e997187 increase benchmark timeout (#12212)
account for compile cache, and it's annoying that job died due to timeout also messes the machine
2025-09-16 14:09:02 -04:00
chenyuandGitHub 84d2d047ea Tensor.pad_to and Tensor.shrink_to (#12210)
most of the time i want this instead of spelling out the args

also add more input validation to shrink
2025-09-16 12:24:55 -04:00
qazalandGitHub 122a50fe8c assert kernel count (#12205) 2025-09-16 14:24:39 +03:00
chenyuandGitHub e555748807 test rangeify const folding (#12200)
* test rangeify const folding

reduce i know how to fix, multi and test_cast_padded tbd

* test_instancenorm_3d is very slow
2025-09-15 20:03:48 -04:00
chenyuandGitHub f732f66709 rangeify test_nn almost pass (#12198)
* rangeify test_nn almost pass

* issue with jit

* flaky
2025-09-15 17:49:20 -04:00
chenyuandGitHub 82e037aad5 ci test.yml updates (#12197)
* ci test.yml updates

move docs together and external_benchmark_schedule to unit

* torch
2025-09-15 17:09:02 -04:00
chenyuandGitHub 146c31586d split RANGEIFY ci (#12196)
one CPU and one CL for speed
2025-09-15 15:41:10 -04:00
chenyuandGitHub df1c183e46 Revert "more llvm intrinsics (#11961)" (#12194)
This reverts commit d01e3d7719.
2025-09-15 13:56:43 -04:00
d01e3d7719 more llvm intrinsics (#11961)
* more llvm intrinsics

* assert nan

* skip test_log_nan on metal

---------

Co-authored-by: b1tg <[email protected]>
2025-09-15 13:05:23 -04:00
nimlgenandGitHub b63bd02969 update runtime docs (#12191) 2025-09-15 17:46:20 +03:00
qazalandGitHub 57e8bf61e8 viz: fix Specificity for rect styling (#12190) 2025-09-15 17:33:37 +03:00
chenyuandGitHub 72e010d816 fix rangeify ci (#12189)
CL=1, and multitensor needs to test with CPU since CL does not support multi in CI
2025-09-15 10:24:57 -04:00
qazalandGitHub f1bd06134d test fuse with RANGEIFY=2 (#12187) 2025-09-15 15:51:23 +03:00
qazalandGitHub ef0ef705fe viz: remove async from event listener (#12186) 2025-09-15 15:08:28 +03:00
qazalandGitHub d8855ec266 viz/serve.py cleanups (#12185)
* don't assign unused variable

* *path to
2025-09-15 13:43:26 +03:00
qazalandGitHub b8a74c1569 cpu: add disassembler err message (#12184)
* cpu: add disassembler err message

* print msg
2025-09-15 13:29:44 +03:00
qazalandGitHub a388d2cb1a remove PROFILE=1 option, it's just VIZ=1 [pr] (#12176)
* remove PROFILE=1 option, it's just VIZ=1 [pr]

* sqtt

* sqtt 2

* return last

* rename
2025-09-15 12:51:50 +03:00
George HotzandGitHub 65397bfdeb set testpath on pytest (#12183) 2025-09-15 16:13:05 +08:00
George HotzandGitHub ae0edc8a67 renumber ranges (#12182)
* enable rangeify const folding

* renumber ranges for kernel deduping
2025-09-15 13:03:39 +08:00
hoovedandGitHub e1fef895b1 don't hardcode weights path (#12171) 2025-09-15 00:33:47 -04:00
hoovedandGitHub 3a9db08b49 download data and ckpts for sd train/eval (#12170) 2025-09-15 00:31:45 -04:00
chenyuandGitHub bdb3afd566 failed test case for symbolic pad (#12179) 2025-09-15 00:25:21 -04:00
George HotzandGitHub 9fcc87761e enable rangeify const folding (#12181) 2025-09-15 12:02:19 +08:00
220 changed files with 13359 additions and 1125 deletions
+11 -11
View File
@@ -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
View File
@@ -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
+4 -2
View File
@@ -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
-2
View File
@@ -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
View File
@@ -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 -1
View File
@@ -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):
+1 -1
View 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
View File
@@ -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):
+2 -1
View File
@@ -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'])
+27
View File
@@ -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:
+63 -1
View File
@@ -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
+23 -2
View File
@@ -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)
+252 -2
View File
@@ -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
+139 -1
View File
@@ -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
View File
@@ -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
@@ -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
+1 -1
View File
@@ -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"
+4 -2
View File
@@ -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"))
+42 -6
View File
@@ -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
View File
@@ -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)
+1 -1
View File
@@ -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 -1
View File
@@ -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
View File
@@ -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
+4 -2
View File
@@ -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
View File
@@ -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 -1
View File
@@ -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
+1 -1
View File
@@ -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):
+11 -1
View File
@@ -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);
+1 -1
View File
@@ -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.
+109
View File
@@ -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())
+74
View File
@@ -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())
+392
View File
@@ -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
+10
View File
@@ -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"
+225
View File
@@ -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 &param, 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 &param, 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"
+3
View File
@@ -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 &param) {
// 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, &param);
}
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 &param) {
#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 &param, 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 &param, 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"
+4
View File
@@ -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");
}
}
}
+213
View File
@@ -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