diff --git a/docs/spec/analysis.md b/docs/spec/analysis.md index 60f48865f..580eda05a 100644 --- a/docs/spec/analysis.md +++ b/docs/spec/analysis.md @@ -166,31 +166,26 @@ target hardware facts. What an occurrence moves is the memory family's answer ```python class ComputeCostMetadata(IRMetadata): - """One Call's logical work, as the authored program states it. - - Attributes: - topologies: attribute; The declared levels, in the order per_unit states them. - flops: attribute; Flop counts per compute DType name. - service: attribute; Result counts per service kind. - """ + """Typed work in logical, expanded, and topology-unit domains.""" topologies: tuple[str, ...] = () flops: Breakdown[int] = Breakdown() - service: Breakdown[int] = Breakdown() + other_ops: Breakdown[int] = Breakdown() ``` + | Field | How it is computed | Reads the target | |---|---|---| | `topologies` | The effective Module topology levels, coarsest first. | No | | `flops` | For a primitive Call, run its registered cost evaluator over operand and result Types as written; the total then multiplies by the enclosing recomputation factor and the number of positions in its execution scope, and each level's share is the same evaluator over Types projected through authored `Split`s at or coarser than that level. For a Function Call, take the callee's summed record and multiply by the call site's factor. | No; projection reads resolved Mesh and effective Module topology extents. | -| `service` | The same evaluator's service counts -- the results it asks a machine for that are not floating point -- totalled and shared the same way. | No; projection reads resolved Mesh and effective Module topology extents. | +| `other_ops` | The evaluator's non-floating-point operation counts (`integer`, `predicate`, `select`, `special`), totalled and shared the same way. These keys map directly to target one-unit operation-throughput keys; target-side service naming is unchanged. | No; projection reads resolved Mesh and effective Module topology extents. | Requesting this family adds one summary line, prefixed by `# `: the Function's own record, stated exactly as a Call's is. The whole program's work is not a second record. ```text -compute-cost flops=:@[,...] service=:@[,...] +compute-cost flops=:@logical,@total,@[,...][;:...] other-ops=:@logical,@total,@[,...][;:...] ``` Every measured Call receives this annotation. Each key pairs the whole quantity @@ -200,8 +195,10 @@ Each reported Call's JSON projection is under its `compute-cost` key: ```text {"topologies": [, ...], - "flops": {: {"total": , "per_unit": [, ...]}}, - "service": {: {"total": , "per_unit": [, ...]}}} + "flops": {: {"logical": , "total": , + "per_unit": [, ...]}}, + "other_ops": {: {"logical": , "total": , + "per_unit": [, ...]}}} ``` - constraints: @@ -234,10 +231,12 @@ class Spread[V]: and a reader who knew to ask. Attributes: - total: attribute; What the whole program asks for. - per_unit: attribute; What one unit of each level holds, in the record's topologies order. + logical: attribute; What the authored operation asks for before replication. + total: attribute; What the whole execution asks for. + per_unit: attribute; What one unit of each level holds, in topology order. """ + logical: V total: V per_unit: tuple[V, ...] = () @@ -541,7 +540,7 @@ Requesting memory adds the Function's own movement, one footprint line, and one line per advisory: ```text -traffic traffic=:r/w@r/w[,...] +traffic traffic=:r/w@total,r/w@[,...][;:...] peak-footprint=:[,:...] advisory="" ``` @@ -560,7 +559,7 @@ is emitted only when asked for ([cli Analyze](./cli.md#analyze)) and is absent from a Function, which has no split: ```text -traffic traffic=:r/w@r/w[,...][ operands=:r/w[,...]] +traffic traffic=:r/w@total,r/w@[,...][;:...] [operands=:r/w[;:...]] ``` Its JSON projection is under the reported value's `traffic` key, with `whole`, diff --git a/docs/tutorial/optimize.ipynb b/docs/tutorial/optimize.ipynb index 6910a0b72..67cb3a31d 100644 --- a/docs/tutorial/optimize.ipynb +++ b/docs/tutorial/optimize.ipynb @@ -54,7 +54,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# traffic traffic=gmem:r71680/w43456@cta:r71680/w43456,thread:r336/w224,rmem:r603496/w488344@cta:r603496/w488344,thread:r2604/w2084\n# roofline ideal-ns=24 bound-by=memory\n" + "text": "# traffic traffic=gmem:r71680/w43456@total,r71680/w43456@cta,r336/w224@thread;rmem:r603496/w488344@total,r603496/w488344@cta,r2604/w2084@thread\n# roofline ideal-ns=24 bound-by=memory\n" } ], "source": [ @@ -94,7 +94,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# traffic traffic=gmem:r43008/w14784@cta:r43008/w14784,thread:r224/w112,rmem:r574824/w459672@cta:r574824/w459672,thread:r2492/w1972\n# roofline ideal-ns=13 bound-by=memory\n" + "text": "# traffic traffic=gmem:r43008/w14784@total,r43008/w14784@cta,r224/w112@thread;rmem:r574824/w459672@total,r574824/w459672@cta,r2492/w1972@thread\n# roofline ideal-ns=13 bound-by=memory\n" } ], "source": [ diff --git a/docs/tutorial/optimize.md b/docs/tutorial/optimize.md index ddac6c2fe..a9aa7790b 100644 --- a/docs/tutorial/optimize.md +++ b/docs/tutorial/optimize.md @@ -88,7 +88,7 @@ grep -E '^# (traffic|roofline) ' Naive.txt ``` ```text -# traffic traffic=gmem:r71680/w43456@cta:r71680/w43456,thread:r336/w224,rmem:r603496/w488344@cta:r603496/w488344,thread:r2604/w2084 +# traffic traffic=gmem:r71680/w43456@total,r71680/w43456@cta,r336/w224@thread;rmem:r603496/w488344@total,r603496/w488344@cta,r2604/w2084@thread # roofline ideal-ns=24 bound-by=memory ``` @@ -132,7 +132,7 @@ grep -E '^# (traffic|roofline) ' Fused.txt ``` ```text -# traffic traffic=gmem:r43008/w14784@cta:r43008/w14784,thread:r224/w112,rmem:r574824/w459672@cta:r574824/w459672,thread:r2492/w1972 +# traffic traffic=gmem:r43008/w14784@total,r43008/w14784@cta,r224/w112@thread;rmem:r574824/w459672@total,r574824/w459672@cta,r2492/w1972@thread # roofline ideal-ns=13 bound-by=memory ``` diff --git a/docs/tutorial/showcase.ipynb b/docs/tutorial/showcase.ipynb index a7aa768fa..ab866d675 100644 --- a/docs/tutorial/showcase.ipynb +++ b/docs/tutorial/showcase.ipynb @@ -90,7 +90,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage0_Naive function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@cta:328896,f32:200448@cta:200448 service=special:1024@cta:1024\n# traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592\n# peak-footprint=gmem:1675788\n# roofline ideal-ns=632 bound-by=memory\n\n v0 = matmul(hidden, w_q, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\"]; compute-cost flops=bf16:131072@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=28 bound-by=memory\n v11 = cache_update(k_cache, cur_pos, write_len, v10) # Tensor[(1, 128, 2, 32), \"bf16\"]; compute-cost; traffic traffic=gmem:r136/w128@cta:r136/w128 operands=0:r0/w0,1:r4/w0,2:r4/w0,3:r128/w0,result:r0/w128; roofline ideal-ns=1 bound-by=memory\n v34 = matmul(v33, w_o, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\"]; compute-cost flops=bf16:131072@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=28 bound-by=memory\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage0_Naive function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,328896@total,328896@cta;f32:200448@logical,200448@total,200448@cta other-ops=special:1024@logical,1024@total,1024@cta\n# traffic traffic=gmem:r2225620/w806592@total,r2225620/w806592@cta\n# peak-footprint=gmem:1675788\n# roofline ideal-ns=632 bound-by=memory\n\n v0 = matmul(hidden, w_q, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta; traffic traffic=gmem:r131584/w512@total,r131584/w512@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=28 bound-by=memory\n v11 = cache_update(k_cache, cur_pos, write_len, v10) # Tensor[(1, 128, 2, 32), \"bf16\"]; compute-cost; traffic traffic=gmem:r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory\n v34 = matmul(v33, w_o, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta; traffic traffic=gmem:r131584/w512@total,r131584/w512@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=28 bound-by=memory\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage0-128.txt\").read_text(encoding=\"utf-8\")\nheader, separator, annotated = report.partition(\"\\n\\n\")\nprint(header.rstrip())\nprint()\nfor needle in (\"matmul(hidden, w_q\", \"cache_update(k_cache\", \"matmul(v33, w_o\"):\n line = next(line for line in annotated.splitlines() if needle in line)\n print(line.rstrip())\n" @@ -137,7 +137,7 @@ { "name": "stdout", "output_type": "stream", - "text": "| `ctx_len` | f32 flops `global@CTA` | traffic `global@CTA` | peak gmem bytes | ideal ns | bound |\n|---:|---:|---|---:|---:|---|\n| 128 | `200448@cta:200448` | `gmem:r2225620/w806592@cta:r2225620/w806592` | 1675788 | 632 | memory |\n| 512 | `799488@cta:799488` | `gmem:r4744660/w3202752@cta:r4744660/w3202752` | 2572812 | 1656 | memory |\n| 1024 | `1598208@cta:1598208` | `gmem:r8103380/w6397632@cta:r8103380/w6397632` | 3768844 | 3022 | memory |\n| 2048 | `3195648@cta:3195648` | `gmem:r14820820/w12787392@cta:r14820820/w12787392` | 6160908 | 5752 | memory |\n| 4096 | `6390528@cta:6390528` | `gmem:r28255700/w25566912@cta:r28255700/w25566912` | 10945036 | 11214 | memory |\n| 8192 | `12780288@cta:12780288` | `gmem:r55125460/w51125952@cta:r55125460/w51125952` | 20513292 | 22136 | memory |\n" + "text": "| `ctx_len` | f32 flops `global@CTA` | traffic `global@CTA` | peak gmem bytes | ideal ns | bound |\n|---:|---:|---|---:|---:|---|\n| 128 | `200448@logical,200448@total,200448@cta` | `gmem:r2225620/w806592@total,r2225620/w806592@cta` | 1675788 | 632 | memory |\n| 512 | `799488@logical,799488@total,799488@cta` | `gmem:r4744660/w3202752@total,r4744660/w3202752@cta` | 2572812 | 1656 | memory |\n| 1024 | `1598208@logical,1598208@total,1598208@cta` | `gmem:r8103380/w6397632@total,r8103380/w6397632@cta` | 3768844 | 3022 | memory |\n| 2048 | `3195648@logical,3195648@total,3195648@cta` | `gmem:r14820820/w12787392@total,r14820820/w12787392@cta` | 6160908 | 5752 | memory |\n| 4096 | `6390528@logical,6390528@total,6390528@cta` | `gmem:r28255700/w25566912@total,r28255700/w25566912@cta` | 10945036 | 11214 | memory |\n| 8192 | `12780288@logical,12780288@total,12780288@cta` | `gmem:r55125460/w51125952@total,r55125460/w51125952@cta` | 20513292 | 22136 | memory |\n" } ], "source": "import re\nfrom pathlib import Path\n\n\ndef metrics(ctx_len):\n report = Path(f\"tutorial-reports/stage0-{ctx_len}.txt\").read_text(encoding=\"utf-8\")\n lines = report.splitlines()\n compute = next(line for line in lines if line.startswith(\"# compute-cost \"))\n traffic = next(line for line in lines if line.startswith(\"# traffic \"))\n peak = next(line for line in lines if line.startswith(\"# peak-footprint=\"))\n roofline = next(line for line in lines if line.startswith(\"# roofline \"))\n f32 = re.search(r\"f32:([^ ]+)\", compute).group(1)\n traffic_value = traffic.removeprefix(\"# traffic traffic=\")\n gmem_peak = re.search(r\"gmem:([^,]+)\", peak).group(1)\n ideal, bound = re.search(r\"ideal-ns=([^ ]+) bound-by=([^ ]+)\", roofline).groups()\n return f32, traffic_value, gmem_peak, ideal, bound\n\n\nprint(\"| `ctx_len` | f32 flops `global@CTA` | traffic `global@CTA` | peak gmem bytes | ideal ns | bound |\")\nprint(\"|---:|---:|---|---:|---:|---|\")\nfor ctx_len in (128, 512, 1024, 2048, 4096, 8192):\n f32, traffic, peak, ideal, bound = metrics(ctx_len)\n print(f\"| {ctx_len} | `{f32}` | `{traffic}` | {peak} | {ideal} | {bound} |\")\n" @@ -188,7 +188,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage2_Sharded function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:2629376@cta:328672,f32:2833728@cta:354216 service=special:14528@cta:1816\n# traffic traffic=gmem:r5563796/w3721856@cta:r3936212/w3721408,smem:r9597248/w9480000@cta:r1199656/w1185000\n# peak-footprint=gmem:3701260,smem:472160\n# roofline ideal-ns=1935 bound-by=memory\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage2_Sharded function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:2833728@logical,2833728@total,354216@cta other-ops=special:14528@logical,14528@total,1816@cta\n# traffic traffic=gmem:r5563796/w3721856@total,r3936212/w3721408@cta;smem:r9597248/w9480000@total,r1199656/w1185000@cta\n# peak-footprint=gmem:3701260;smem:472160\n# roofline ideal-ns=1935 bound-by=memory\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage2-1816.txt\").read_text(encoding=\"utf-8\")\nprint(report.partition(\"\\n\\n\")[0].rstrip())\n" @@ -274,7 +274,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# compute-cost flops=bf16:328896@cta:328896,f32:200448@cta:200448 service=special:1024@cta:1024\n# traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592\n# peak-footprint=gmem:1675788\n# roofline ideal-ns=632 bound-by=memory\n" + "text": "# compute-cost flops=bf16:328896@logical,328896@total,328896@cta;f32:200448@logical,200448@total,200448@cta other-ops=special:1024@logical,1024@total,1024@cta\n# traffic traffic=gmem:r2225620/w806592@total,r2225620/w806592@cta\n# peak-footprint=gmem:1675788\n# roofline ideal-ns=632 bound-by=memory\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage0-128-summary.txt\").read_text(encoding=\"utf-8\")\nfor line in report.splitlines():\n if line.startswith((\"# compute-cost \", \"# traffic \", \"# peak-footprint=\", \"# roofline \")):\n print(line)\n" @@ -309,7 +309,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# compute-cost flops=bf16:2629376@cta:328672,f32:200448@cta:25056 service=special:1024@cta:128\n# traffic traffic=gmem:r1674644/w264832@cta:r1559508/w264384,smem:r684608/w675392@cta:r85576/w84424\n# peak-footprint=gmem:1540620,smem:33280\n# roofline ideal-ns=405 bound-by=memory\n" + "text": "# compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:200448@logical,200448@total,25056@cta other-ops=special:1024@logical,1024@total,128@cta\n# traffic traffic=gmem:r1674644/w264832@total,r1559508/w264384@cta;smem:r684608/w675392@total,r85576/w84424@cta\n# peak-footprint=gmem:1540620;smem:33280\n# roofline ideal-ns=405 bound-by=memory\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage2-128-summary.txt\").read_text(encoding=\"utf-8\")\nfor line in report.splitlines():\n if line.startswith((\"# compute-cost \", \"# traffic \", \"# peak-footprint=\", \"# roofline \")):\n print(line)\n" @@ -360,7 +360,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage3_Fused function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:4392896@cta:328672,f32:6418944@cta:200592 service=integer:288@cta:9,special:33152@cta:1036\n# traffic traffic=gmem:r3476884/w4196992@cta:r2558932/w4196544,rmem:r656/w72@cta:r656/w72,smem:r5839296/w5662784@cta:r682464/w672676\n# peak-footprint=gmem:5047436,rmem:16,smem:33408\n# roofline ideal-ns=1599 bound-by=memory\n\n v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), \"bf16\"]; compute-cost; traffic traffic=gmem:r136/w128@cta:r136/w128 operands=0:r0/w0,1:r4/w0,2:r4/w0,3:r128/w0,result:r0/w128; roofline ideal-ns=1 bound-by=memory\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage3_Fused function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,4392896@total,328672@cta;f32:3239808@logical,6418944@total,200592@cta other-ops=integer:9@logical,288@total,9@cta;special:33056@logical,33152@total,1036@cta\n# traffic traffic=gmem:r3476884/w4196992@total,r2558932/w4196544@cta;rmem:r656/w72@total,r656/w72@cta;smem:r5839296/w5662784@total,r682464/w672676@cta\n# peak-footprint=gmem:5047436;rmem:16;smem:33408\n# roofline ideal-ns=1599 bound-by=memory\n\n v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), \"bf16\"]; compute-cost; traffic traffic=gmem:r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage3-4096.txt\").read_text(encoding=\"utf-8\")\nheader, separator, annotated = report.partition(\"\\n\\n\")\nprint(header.rstrip())\nprint()\nprint(next(line.rstrip() for line in annotated.splitlines() if \"cache_update(k_cache\" in line))\n" @@ -411,7 +411,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage4_WeightPrepared function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:337408@cta:42176,f32:51124224@cta:6390528 service=special:262144@cta:32768\n# traffic traffic=gmem:r28254676/w25566912@cta:r27967956/w25565792,smem:r331008/w329984@cta:r43168/w42144\n# peak-footprint=gmem:10945036,smem:16960\n# roofline ideal-ns=11213 bound-by=memory\n\n v1 = reshard(w_q, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n v2 = matmul(v0, v1, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), \"smem\"]; compute-cost flops=bf16:131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=1 bound-by=compute\n\n v42 = reshard(w_o, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n v43 = matmul(v41, v42, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), \"smem\"]; compute-cost flops=bf16:131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=1 bound-by=compute\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage4_WeightPrepared function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,337408@total,42176@cta;f32:6390528@logical,51124224@total,6390528@cta other-ops=special:32768@logical,262144@total,32768@cta\n# traffic traffic=gmem:r28254676/w25566912@total,r27967956/w25565792@cta;smem:r331008/w329984@total,r43168/w42144@cta\n# peak-footprint=gmem:10945036;smem:16960\n# roofline ideal-ns=11213 bound-by=memory\n\n v1 = reshard(w_q, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@total,r16384/w0@cta;smem:r0/w131072@total,r0/w16384@cta operands=0:r131072/w0;result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n v2 = matmul(v0, v1, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), \"smem\"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta; traffic traffic=smem:r131584/w512@total,r16896/w64@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=1 bound-by=compute\n\n v42 = reshard(w_o, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), \"bf16\", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), \"smem\"]; compute-cost; traffic traffic=gmem:r131072/w0@total,r16384/w0@cta;smem:r0/w131072@total,r0/w16384@cta operands=0:r131072/w0;result:r0/w131072; roofline ideal-ns=28 bound-by=memory\n v43 = matmul(v41, v42, a_layout=\"MK\", b_layout=\"KN\") # Tensor[(1, 1, 256), \"bf16\", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), \"smem\"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta; traffic traffic=smem:r131584/w512@total,r16896/w64@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=1 bound-by=compute\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage4-4096.txt\").read_text(encoding=\"utf-8\")\nheader, separator, annotated = report.partition(\"\\n\\n\")\nprint(header.rstrip())\nlines = annotated.splitlines()\nfor needle in (\"reshard(w_q\", \"reshard(w_o\"):\n start = next(index for index, line in enumerate(lines) if needle in line)\n end = start\n while end + 1 < len(lines):\n end += 1\n if end > start and \" # \" in lines[end]:\n break\n print()\n print(\"\\n\".join(line.rstrip() for line in lines[start : end + 1]))\n" @@ -462,7 +462,7 @@ { "name": "stdout", "output_type": "stream", - "text": "# analysis target=nvidia.h200_sxm module=Stage5_CachePrepared function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:1246400@cta:328672,f32:6410280@cta:801285 service=integer:256@cta:32,special:33040@cta:4130\n# traffic traffic=gmem:r7672476/w4198272@cta:r4001116/w4197824,rmem:r2560/w0@cta:r2560/w0,smem:r21788704/w21486432@cta:r2723588/w2685804\n# peak-footprint=gmem:2950412,rmem:0,smem:33536\n# roofline ideal-ns=2474 bound-by=memory\n\n v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), \"bf16\"]; compute-cost; traffic traffic=rmem:r32/w0@cta:r32/w0 operands=0:r0/w0,1:r32/w0,result:r0/w0; roofline\n v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), \"bf16\"]; compute-cost; traffic traffic=gmem:r136/w128@cta:r136/w128 operands=0:r0/w0,1:r4/w0,2:r4/w0,3:r128/w0,result:r0/w128; roofline ideal-ns=1 bound-by=memory\n" + "text": "# analysis target=nvidia.h200_sxm module=Stage5_CachePrepared function=gqa_decode topology=cta\n# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline\n# compute-cost flops=bf16:328896@logical,1246400@total,328672@cta;f32:6410280@logical,6410280@total,801285@cta other-ops=integer:32@logical,256@total,32@cta;special:33040@logical,33040@total,4130@cta\n# traffic traffic=gmem:r7672476/w4198272@total,r4001116/w4197824@cta;rmem:r2560/w0@total,r2560/w0@cta;smem:r21788704/w21486432@total,r2723588/w2685804@cta\n# peak-footprint=gmem:2950412;rmem:0;smem:33536\n# roofline ideal-ns=2474 bound-by=memory\n\n v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), \"bf16\"]; compute-cost; traffic traffic=rmem:r32/w0@total,r32/w0@cta operands=0:r0/w0;1:r32/w0;result:r0/w0; roofline\n v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), \"bf16\"]; compute-cost; traffic traffic=gmem:r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory\n" } ], "source": "from pathlib import Path\n\nreport = Path(\"tutorial-reports/stage5-4096.txt\").read_text(encoding=\"utf-8\")\nheader, separator, annotated = report.partition(\"\\n\\n\")\nprint(header.rstrip())\nprint()\nfor needle in (\"slice(k_cache\", \"cache_update(k_cache\"):\n print(next(line.rstrip() for line in annotated.splitlines() if needle in line))\n" diff --git a/docs/tutorial/showcase.md b/docs/tutorial/showcase.md index 45bb5856d..2305af926 100644 --- a/docs/tutorial/showcase.md +++ b/docs/tutorial/showcase.md @@ -197,14 +197,14 @@ for needle in ("matmul(hidden, w_q", "cache_update(k_cache", "matmul(v33, w_o"): ```text # analysis target=nvidia.h200_sxm module=Stage0_Naive function=gqa_decode topology=cta # selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline -# compute-cost flops=bf16:328896@cta:328896,f32:200448@cta:200448 service=special:1024@cta:1024 -# traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592 +# compute-cost flops=bf16:328896@logical,328896@total,328896@cta;f32:200448@logical,200448@total,200448@cta other-ops=special:1024@logical,1024@total,1024@cta +# traffic traffic=gmem:r2225620/w806592@total,r2225620/w806592@cta # peak-footprint=gmem:1675788 # roofline ideal-ns=632 bound-by=memory - v0 = matmul(hidden, w_q, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16"]; compute-cost flops=bf16:131072@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=28 bound-by=memory - v11 = cache_update(k_cache, cur_pos, write_len, v10) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@cta:r136/w128 operands=0:r0/w0,1:r4/w0,2:r4/w0,3:r128/w0,result:r0/w128; roofline ideal-ns=1 bound-by=memory - v34 = matmul(v33, w_o, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16"]; compute-cost flops=bf16:131072@cta:131072; traffic traffic=gmem:r131584/w512@cta:r131584/w512 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=28 bound-by=memory + v0 = matmul(hidden, w_q, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta; traffic traffic=gmem:r131584/w512@total,r131584/w512@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=28 bound-by=memory + v11 = cache_update(k_cache, cur_pos, write_len, v10) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory + v34 = matmul(v33, w_o, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta; traffic traffic=gmem:r131584/w512@total,r131584/w512@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=28 bound-by=memory ``` The number before `@` is global work or traffic. After it comes one share per @@ -259,12 +259,12 @@ for ctx_len in (128, 512, 1024, 2048, 4096, 8192): | `ctx_len` | f32 flops `global@CTA` | traffic `global@CTA` | peak gmem bytes | ideal ns | bound | |---:|---:|---|---:|---:|---| -| 128 | `200448@cta:200448` | `gmem:r2225620/w806592@cta:r2225620/w806592` | 1675788 | 632 | memory | -| 512 | `799488@cta:799488` | `gmem:r4744660/w3202752@cta:r4744660/w3202752` | 2572812 | 1656 | memory | -| 1024 | `1598208@cta:1598208` | `gmem:r8103380/w6397632@cta:r8103380/w6397632` | 3768844 | 3022 | memory | -| 2048 | `3195648@cta:3195648` | `gmem:r14820820/w12787392@cta:r14820820/w12787392` | 6160908 | 5752 | memory | -| 4096 | `6390528@cta:6390528` | `gmem:r28255700/w25566912@cta:r28255700/w25566912` | 10945036 | 11214 | memory | -| 8192 | `12780288@cta:12780288` | `gmem:r55125460/w51125952@cta:r55125460/w51125952` | 20513292 | 22136 | memory | +| 128 | `200448@logical,200448@total,200448@cta` | `gmem:r2225620/w806592@total,r2225620/w806592@cta` | 1675788 | 632 | memory | +| 512 | `799488@logical,799488@total,799488@cta` | `gmem:r4744660/w3202752@total,r4744660/w3202752@cta` | 2572812 | 1656 | memory | +| 1024 | `1598208@logical,1598208@total,1598208@cta` | `gmem:r8103380/w6397632@total,r8103380/w6397632@cta` | 3768844 | 3022 | memory | +| 2048 | `3195648@logical,3195648@total,3195648@cta` | `gmem:r14820820/w12787392@total,r14820820/w12787392@cta` | 6160908 | 5752 | memory | +| 4096 | `6390528@logical,6390528@total,6390528@cta` | `gmem:r28255700/w25566912@total,r28255700/w25566912@cta` | 10945036 | 11214 | memory | +| 8192 | `12780288@logical,12780288@total,12780288@cta` | `gmem:r55125460/w51125952@total,r55125460/w51125952@cta` | 20513292 | 22136 | memory | The table says: @@ -426,9 +426,9 @@ print(report.partition("\n\n")[0].rstrip()) ```text # analysis target=nvidia.h200_sxm module=Stage2_Sharded function=gqa_decode topology=cta # selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline -# compute-cost flops=bf16:2629376@cta:328672,f32:2833728@cta:354216 service=special:14528@cta:1816 -# traffic traffic=gmem:r5563796/w3721856@cta:r3936212/w3721408,smem:r9597248/w9480000@cta:r1199656/w1185000 -# peak-footprint=gmem:3701260,smem:472160 +# compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:2833728@logical,2833728@total,354216@cta other-ops=special:14528@logical,14528@total,1816@cta +# traffic traffic=gmem:r5563796/w3721856@total,r3936212/w3721408@cta;smem:r9597248/w9480000@total,r1199656/w1185000@cta +# peak-footprint=gmem:3701260;smem:472160 # roofline ideal-ns=1935 bound-by=memory ``` @@ -554,8 +554,8 @@ for line in report.splitlines(): ``` ```text -# compute-cost flops=bf16:328896@cta:328896,f32:200448@cta:200448 service=special:1024@cta:1024 -# traffic traffic=gmem:r2225620/w806592@cta:r2225620/w806592 +# compute-cost flops=bf16:328896@logical,328896@total,328896@cta;f32:200448@logical,200448@total,200448@cta other-ops=special:1024@logical,1024@total,1024@cta +# traffic traffic=gmem:r2225620/w806592@total,r2225620/w806592@cta # peak-footprint=gmem:1675788 # roofline ideal-ns=632 bound-by=memory ``` @@ -581,9 +581,9 @@ for line in report.splitlines(): ``` ```text -# compute-cost flops=bf16:2629376@cta:328672,f32:200448@cta:25056 service=special:1024@cta:128 -# traffic traffic=gmem:r1674644/w264832@cta:r1559508/w264384,smem:r684608/w675392@cta:r85576/w84424 -# peak-footprint=gmem:1540620,smem:33280 +# compute-cost flops=bf16:328896@logical,2629376@total,328672@cta;f32:200448@logical,200448@total,25056@cta other-ops=special:1024@logical,1024@total,128@cta +# traffic traffic=gmem:r1674644/w264832@total,r1559508/w264384@cta;smem:r684608/w675392@total,r85576/w84424@cta +# peak-footprint=gmem:1540620;smem:33280 # roofline ideal-ns=405 bound-by=memory ``` @@ -728,12 +728,12 @@ print(next(line.rstrip() for line in annotated.splitlines() if "cache_update(k_c ```text # analysis target=nvidia.h200_sxm module=Stage3_Fused function=gqa_decode topology=cta # selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline -# compute-cost flops=bf16:4392896@cta:328672,f32:6418944@cta:200592 service=integer:288@cta:9,special:33152@cta:1036 -# traffic traffic=gmem:r3476884/w4196992@cta:r2558932/w4196544,rmem:r656/w72@cta:r656/w72,smem:r5839296/w5662784@cta:r682464/w672676 -# peak-footprint=gmem:5047436,rmem:16,smem:33408 +# compute-cost flops=bf16:328896@logical,4392896@total,328672@cta;f32:3239808@logical,6418944@total,200592@cta other-ops=integer:9@logical,288@total,9@cta;special:33056@logical,33152@total,1036@cta +# traffic traffic=gmem:r3476884/w4196992@total,r2558932/w4196544@cta;rmem:r656/w72@total,r656/w72@cta;smem:r5839296/w5662784@total,r682464/w672676@cta +# peak-footprint=gmem:5047436;rmem:16;smem:33408 # roofline ideal-ns=1599 bound-by=memory - v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@cta:r136/w128 operands=0:r0/w0,1:r4/w0,2:r4/w0,3:r128/w0,result:r0/w128; roofline ideal-ns=1 bound-by=memory + v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory ``` The embedded `Stage3_Fused` program is the split-K example for this page. @@ -865,16 +865,16 @@ for needle in ("reshard(w_q", "reshard(w_o"): ```text # analysis target=nvidia.h200_sxm module=Stage4_WeightPrepared function=gqa_decode topology=cta # selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline -# compute-cost flops=bf16:337408@cta:42176,f32:51124224@cta:6390528 service=special:262144@cta:32768 -# traffic traffic=gmem:r28254676/w25566912@cta:r27967956/w25565792,smem:r331008/w329984@cta:r43168/w42144 -# peak-footprint=gmem:10945036,smem:16960 +# compute-cost flops=bf16:328896@logical,337408@total,42176@cta;f32:6390528@logical,51124224@total,6390528@cta other-ops=special:32768@logical,262144@total,32768@cta +# traffic traffic=gmem:r28254676/w25566912@total,r27967956/w25565792@cta;smem:r331008/w329984@total,r43168/w42144@cta +# peak-footprint=gmem:10945036;smem:16960 # roofline ideal-ns=11213 bound-by=memory - v1 = reshard(w_q, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory - v2 = matmul(v0, v1, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), "smem"]; compute-cost flops=bf16:131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=1 bound-by=compute + v1 = reshard(w_q, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; traffic traffic=gmem:r131072/w0@total,r16384/w0@cta;smem:r0/w131072@total,r0/w16384@cta operands=0:r131072/w0;result:r0/w131072; roofline ideal-ns=28 bound-by=memory + v2 = matmul(v0, v1, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), "smem"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta; traffic traffic=smem:r131584/w512@total,r16896/w64@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=1 bound-by=compute - v42 = reshard(w_o, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; traffic traffic=gmem:r131072/w0@cta:r16384/w0,smem:r0/w131072@cta:r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory - v43 = matmul(v41, v42, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), "smem"]; compute-cost flops=bf16:131072@cta:16384; traffic traffic=smem:r131584/w512@cta:r16896/w64 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=1 bound-by=compute + v42 = reshard(w_o, layout=(1, 256, 8 @ mesh.head, 32), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; traffic traffic=gmem:r131072/w0@total,r16384/w0@cta;smem:r0/w131072@total,r0/w16384@cta operands=0:r131072/w0;result:r0/w131072; roofline ideal-ns=28 bound-by=memory + v43 = matmul(v41, v42, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), "smem"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta; traffic traffic=smem:r131584/w512@total,r16896/w64@cta operands=0:r512/w0;1:r131072/w0;result:r0/w512; roofline ideal-ns=1 bound-by=compute ``` ## 6. Stream the KV cache @@ -1017,13 +1017,13 @@ for needle in ("slice(k_cache", "cache_update(k_cache"): ```text # analysis target=nvidia.h200_sxm module=Stage5_CachePrepared function=gqa_decode topology=cta # selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline -# compute-cost flops=bf16:1246400@cta:328672,f32:6410280@cta:801285 service=integer:256@cta:32,special:33040@cta:4130 -# traffic traffic=gmem:r7672476/w4198272@cta:r4001116/w4197824,rmem:r2560/w0@cta:r2560/w0,smem:r21788704/w21486432@cta:r2723588/w2685804 -# peak-footprint=gmem:2950412,rmem:0,smem:33536 +# compute-cost flops=bf16:328896@logical,1246400@total,328672@cta;f32:6410280@logical,6410280@total,801285@cta other-ops=integer:32@logical,256@total,32@cta;special:33040@logical,33040@total,4130@cta +# traffic traffic=gmem:r7672476/w4198272@total,r4001116/w4197824@cta;rmem:r2560/w0@total,r2560/w0@cta;smem:r21788704/w21486432@total,r2723588/w2685804@cta +# peak-footprint=gmem:2950412;rmem:0;smem:33536 # roofline ideal-ns=2474 bound-by=memory - v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost; traffic traffic=rmem:r32/w0@cta:r32/w0 operands=0:r0/w0,1:r32/w0,result:r0/w0; roofline - v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@cta:r136/w128 operands=0:r0/w0,1:r4/w0,2:r4/w0,3:r128/w0,result:r0/w128; roofline ideal-ns=1 bound-by=memory + v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost; traffic traffic=rmem:r32/w0@total,r32/w0@cta operands=0:r0/w0;1:r32/w0;result:r0/w0; roofline + v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@total,r136/w128@cta operands=0:r0/w0;1:r4/w0;2:r4/w0;3:r128/w0;result:r0/w128; roofline ideal-ns=1 bound-by=memory ``` ```text diff --git a/src/tilefoundry/analysis/api.py b/src/tilefoundry/analysis/api.py index f61e5df01..e1ef0ea25 100644 --- a/src/tilefoundry/analysis/api.py +++ b/src/tilefoundry/analysis/api.py @@ -46,6 +46,11 @@ class AnalysisResult: executed: tuple[str, ...] metadata_types: tuple[type[IRMetadata], ...] + @property + def level(self) -> str | None: + """The selected topology level.""" + return self.topology_level + def _algorithm(target: Target, selector: str, *, root: str) -> Analyzer: """The service selected by the resolved Target for *selector*.""" @@ -107,6 +112,7 @@ def analyze( *, analysis: str | Sequence[str], topology_level: str | None = None, + level: str | None = None, options: object | None = None, dims: "Mapping[str, int] | None" = None, ) -> AnalysisResult: @@ -118,6 +124,9 @@ def analyze( variant owned by the module, and the result identifies the concrete inlined view that received records. """ + if topology_level is not None and level is not None and topology_level != level: + raise ValueError("analyze: topology_level and level disagree") + topology_level = topology_level if topology_level is not None else level if not isinstance(module, Module): raise TypeError( f"analyze: expected a Module, got {type(module).__name__}. A Function " diff --git a/src/tilefoundry/analysis/compute_cost.py b/src/tilefoundry/analysis/compute_cost.py index bd4f73609..c8907c334 100644 --- a/src/tilefoundry/analysis/compute_cost.py +++ b/src/tilefoundry/analysis/compute_cost.py @@ -1,11 +1,4 @@ -"""How much work the authored program asks for. - -This family reads the program and nothing else. Flops and typed service come -from each op's registered cost evaluator, so the record it leaves is the same on -every backend. What that work moves is the memory family's half of the same -declaration, and what it costs in time is a separate question again, asked -against a target's rates. -""" +"""Count typed floating-point and other operations in authored HIR.""" from __future__ import annotations @@ -20,11 +13,7 @@ from tilefoundry.ir.types.shard import Mesh, composed, topology_axes from tilefoundry.ir.types.shard.mesh import _positions_layout from tilefoundry.ir.visitor import ExprVisitor -from tilefoundry.visitor_registry.contexts import ( - CostContext, - FunctionScope, - TrafficBytes, -) +from tilefoundry.visitor_registry.contexts import CostContext, FunctionScope, TrafficBytes from tilefoundry.visitor_registry.visitors import CostEvaluator from .errors import AnalysisError @@ -35,6 +24,14 @@ SELECTOR = "compute-cost" +def _counts(values: dict) -> tuple[tuple[str, int], ...]: + return tuple(sorted((getattr(kind, "name", kind), value) for kind, value in values.items())) + + +def _at(held: Breakdown[int], topologies: tuple[str, ...], level: str | None): + return tuple(shares(held, topologies, level).items()) + + def _is_structural_occurrence( cost: ComputeCostMetadata, moved: "TrafficMetadata | None" = None, @@ -42,18 +39,10 @@ def _is_structural_occurrence( unit: str, bandwidth_level: str | None = None, ) -> bool: - """Whether an occurrence asks for nothing this model puts on a clock. - - Only what could take time is counted: the flops, the typed service, and the - bytes at the one level a bandwidth is published for. Movement at any other - level is still movement and still recorded -- what it is not is work this - model can lay on a timeline, so it neither earns a duration nor asks for a - placement to be laid at. Having moved bytes and having timed work are - different questions, and this is the second one. - """ + """Whether an occurrence asks for nothing this model puts on a clock.""" return ( - all(not value for _name, value in shares(cost.flops, cost.topologies, unit).items()) - and all(not value for _kind, value in shares(cost.service, cost.topologies, unit).items()) + all(not value for _name, value in _at(cost.flops, cost.topologies, unit)) + and all(not value for _kind, value in _at(cost.other_ops, cost.topologies, unit)) and not ( _bytes(moved.storage, moved.topologies, bandwidth_level, unit) if moved is not None and bandwidth_level is not None @@ -68,32 +57,26 @@ def _local_duration_ns( services: PerformanceServiceFacts, *, moved: "TrafficMetadata | None" = None, - topology_level: str, + topology_level: str | None = None, + level: str | None = None, scale: int = 1, ) -> int: - """Price one occurrence's projected work against one unit's throughputs. - - Compute, movement and what leaves the unit overlap, so the duration is - whichever takes longest; the same bytes can spend two resources and still - take one span of time. Work with no stated throughput is refused rather - than priced at nothing. Movement at a level the target publishes no - bandwidth for is a different case -- stated and left untimed, because a - rate nobody published is not one this may invent. - """ + """Price one occurrence's projected work against one unit's throughputs.""" + topology_level = topology_level if topology_level is not None else level + if topology_level is None: + raise AnalysisError("performance: no topology level was selected") if services.unit != topology_level: raise AnalysisError( f"performance: selected topology level {topology_level!r}, but the " f"target's one-unit throughputs are stated for {services.unit!r}" ) - if _is_structural_occurrence( cost, moved, unit=topology_level, bandwidth_level=facts.bandwidth_level ): return 0 compute_ns = 0 - for name, share in shares(cost.flops, cost.topologies, topology_level).items(): - value = share + for name, value in _at(cost.flops, cost.topologies, topology_level): if not value: continue dtype = getattr(DType, name, None) @@ -107,8 +90,7 @@ def _local_duration_ns( ) compute_ns += -(-(value * scale * 1_000_000_000) // throughput) - for kind, share in shares(cost.service, cost.topologies, topology_level).items(): - value = share + for kind, value in _at(cost.other_ops, cost.topologies, topology_level): if not value: continue throughput = services.ops(kind) @@ -147,48 +129,53 @@ def _local_duration_ns( return max(compute_ns, memory_ns, link_ns) -def _named(flops: dict) -> dict[str, int]: - """The flop counts keyed by dtype name, as a record states them.""" - return {dtype.name: value for dtype, value in flops.items()} - - def _call_cost_record( expr: Call, - locals_by_unit: "dict[str, CostContext]", - positions_by_unit: "dict[str, int]", - asked: "str | None" = None, + locals_by_unit: dict[str, CostContext], + positions_by_unit: dict[str, int], + whole: CostContext, + asked: str | None = None, ) -> ComputeCostMetadata: - """Measure the work one Call asks for, without attaching the record. - - Work only: what an occurrence moves is the memory family's answer. One - unit's work is measured at every declared level rather than at one chosen - for the reader, and global work is any of them repeated over the positions - executing this scope -- the product is the same whichever level states it. - """ - flops_by_unit: list[dict[str, int]] = [] - service_by_unit: list[dict[str, int]] = [] - whole_flops: dict[str, int] = {} - whole_service: dict[str, int] = {} + """Measure one Call in logical, expanded, and every topology-unit domain.""" + flops_by_unit: list[tuple[tuple[str, int], ...]] = [] + other_ops_by_unit: list[tuple[tuple[str, int], ...]] = [] + total_flops: tuple[tuple[str, int], ...] = () + total_other_ops: tuple[tuple[str, int], ...] = () + try: + logical = CostEvaluator().visit(expr, whole) + except (ValueError, VerifyError) as error: + raise AnalysisError(str(error)) from None for unit, local in locals_by_unit.items(): try: cost = CostEvaluator().visit(expr, local) except (ValueError, VerifyError) as error: raise AnalysisError(str(error)) from None - flops_by_unit.append(_named(cost.flops)) - service_by_unit.append(dict(cost.service)) + unit_flops = _counts(cost.flops) + unit_other_ops = _counts(cost.service) + flops_by_unit.append(unit_flops) + other_ops_by_unit.append(unit_other_ops) if unit == (asked or next(iter(locals_by_unit))): repeats = positions_by_unit[unit] - whole_flops = {name: value * repeats for name, value in _named(cost.flops).items()} - whole_service = {kind: value * repeats for kind, value in cost.service.items()} + total_flops = tuple((name, value * repeats) for name, value in unit_flops) + total_other_ops = tuple((kind, value * repeats) for kind, value in unit_other_ops) return ComputeCostMetadata( topologies=tuple(locals_by_unit), - flops=breakdown(whole_flops, flops_by_unit, 0), - service=breakdown(whole_service, service_by_unit, 0), + flops=breakdown( + dict(total_flops), + tuple(dict(values) for values in flops_by_unit), + 0, + logical=dict(_counts(logical.flops)), + ), + other_ops=breakdown( + dict(total_other_ops), + tuple(dict(values) for values in other_ops_by_unit), + 0, + logical=dict(_counts(logical.service)), + ), ) def _scope_position_count(mesh: Mesh, topology_level: str | None, topologies: tuple) -> int: - """Count positions at or above the selected level within *mesh*.""" if topology_level is None: return 1 declared = {topology.name: index for index, topology in enumerate(topologies)} @@ -209,49 +196,49 @@ def _scope_position_count(mesh: Mesh, topology_level: str | None, topologies: tu def _bytes( - held: "Breakdown[TrafficBytes]", + held: Breakdown[TrafficBytes], topologies: tuple[str, ...], kind: str, - topology_level: "str | None", + topology_level: str | None, ) -> int: - """One kind's bytes for one unit of *topology_level*, read and written together.""" moved = shares(held, topologies, topology_level).get(kind) return moved.total_bytes if moved is not None else 0 -def _accumulate( - flops: dict[str, int], - service: dict[str, int], - by_unit: "dict[str, dict[str, dict[str, int]]]", - record: ComputeCostMetadata, - trips: int, -) -> None: - for kind, spread in record.flops.kinds: - flops[kind] = flops.get(kind, 0) + spread.total * trips - for kind, spread in record.service.kinds: - service[kind] = service.get(kind, 0) + spread.total * trips +def _add(target: dict[str, int], values, trips: int) -> None: + for name, value in values: + target[name] = target.get(name, 0) + value * trips + + +def _domain_values(held: Breakdown[int], domain: str) -> tuple[tuple[str, int], ...]: + return tuple((name, getattr(spread, domain)) for name, spread in held.kinds) + + +def _accumulate(ctx: "ComputeCostContext", record: ComputeCostMetadata, trips: int) -> None: + _add(ctx.flops_logical, _domain_values(record.flops, "logical"), trips) + _add(ctx.flops, _domain_values(record.flops, "total"), trips) + _add(ctx.other_ops_logical, _domain_values(record.other_ops, "logical"), trips) + _add(ctx.other_ops, _domain_values(record.other_ops, "total"), trips) for index, unit in enumerate(record.topologies): - held = by_unit.setdefault(unit, {"flops": {}, "service": {}}) - for kind, spread in record.flops.kinds: - held["flops"][kind] = held["flops"].get(kind, 0) + spread.at(index) * trips - for kind, spread in record.service.kinds: - held["service"][kind] = held["service"].get(kind, 0) + spread.at(index) * trips + held = ctx.by_unit.setdefault(unit, {"flops": {}, "other_ops": {}}) + _add(held["flops"], ((name, spread.at(index)) for name, spread in record.flops.kinds), trips) + _add( + held["other_ops"], + ((name, spread.at(index)) for name, spread in record.other_ops.kinds), + trips, + ) @dataclass class ComputeCostContext(AnalyzeContext): - """State carried through the compute-cost expression walk. - - ``executing_positions`` is how many positions of each declared level the - enclosing scope holds; it is the multiplier that turns one unit's work at - that level into total replicated work. - """ - locals_by_unit: dict[str, CostContext] = field(default_factory=dict) + whole: CostContext | None = None current_mesh: Mesh | None = None executing_positions: dict[str, int] = field(default_factory=dict) flops: dict[str, int] = field(default_factory=dict) - service: dict[str, int] = field(default_factory=dict) + flops_logical: dict[str, int] = field(default_factory=dict) + other_ops: dict[str, int] = field(default_factory=dict) + other_ops_logical: dict[str, int] = field(default_factory=dict) by_unit: dict[str, dict[str, dict[str, int]]] = field(default_factory=dict) call_count: list[int] = field(default_factory=lambda: [0]) @@ -260,7 +247,6 @@ class ComputeCostVisitor(ExprVisitor[None]): """Attach per-Call work and accumulate multiplicity-aware totals.""" def visit_MeshRegion(self, expr: MeshRegion, ctx: ComputeCostContext) -> None: - """Carry the region's execution multiplicity into each contained Call.""" for arg in expr.args: self.visit(arg, ctx) mesh = composed((ctx.current_mesh, expr.mesh)) if ctx.current_mesh else expr.mesh @@ -268,10 +254,7 @@ def visit_MeshRegion(self, expr: MeshRegion, ctx: ComputeCostContext) -> None: positions = { unit: _scope_position_count(mesh, unit, topologies) for unit in ctx.locals_by_unit } - self.visit( - expr.body, - replace(ctx, executing_positions=positions, current_mesh=mesh), - ) + self.visit(expr.body, replace(ctx, executing_positions=positions, current_mesh=mesh)) def visit_LoopRegion(self, expr: LoopRegion, ctx: ComputeCostContext) -> None: child = next(item for item in ctx.current.children if item.owner is expr) @@ -288,10 +271,10 @@ def default_visit_leaf( if not isinstance(expr, Call): return ctx.call_count[0] += 1 - if not ctx.locals_by_unit: + if not ctx.locals_by_unit or ctx.whole is None: raise AnalysisError("compute-cost: visitor context is missing its cost context") record = _call_cost_record( - expr, ctx.locals_by_unit, ctx.executing_positions, ctx.topology_level + expr, ctx.locals_by_unit, ctx.executing_positions, ctx.whole, ctx.topology_level ) attach(expr, record) owner = ctx.current if id(expr) in ctx.current.accesses["narrow"] else ctx.root @@ -301,13 +284,10 @@ def default_visit_leaf( if cursor.is_variant(expr): repeats *= max(1, cursor.trips()) cursor = cursor.parent - _accumulate(ctx.flops, ctx.service, ctx.by_unit, record, repeats) + _accumulate(ctx, record, repeats) -def analyze_compute_cost( - function: Function, - context: AnalyzeContext, -) -> None: +def analyze_compute_cost(function: Function, context: AnalyzeContext) -> None: """Attach one-trip work per Call and multiplicity-aware totals per Function.""" module, topology_level = context.module, context.topology_level topologies = module.effective_topologies() @@ -326,23 +306,27 @@ def analyze_compute_cost( root=context.root, current=context.current, locals_by_unit=locals_by_unit, + whole=CostContext(scope=scope), executing_positions=dict.fromkeys(locals_by_unit, 1), ) ComputeCostVisitor().visit(function.body, cost_context) if cost_context.call_count[0] > 0: + units = tuple(cost_context.by_unit) attach( function, ComputeCostMetadata( - topologies=tuple(cost_context.by_unit), + topologies=units, flops=breakdown( - dict(cost_context.flops), - [dict(held["flops"]) for held in cost_context.by_unit.values()], + cost_context.flops, + tuple(cost_context.by_unit[u]["flops"] for u in units), 0, + logical=cost_context.flops_logical, ), - service=breakdown( - dict(cost_context.service), - [dict(held["service"]) for held in cost_context.by_unit.values()], + other_ops=breakdown( + cost_context.other_ops, + tuple(cost_context.by_unit[u]["other_ops"] for u in units), 0, + logical=cost_context.other_ops_logical, ), ), ) diff --git a/src/tilefoundry/analysis/facts.py b/src/tilefoundry/analysis/facts.py index 6967cf6ec..b1756cd3f 100644 --- a/src/tilefoundry/analysis/facts.py +++ b/src/tilefoundry/analysis/facts.py @@ -89,15 +89,11 @@ class MemoryHierarchyFacts: def explicit(self, name: str) -> ExplicitMemoryLevelFacts | None: """The explicit level called *name*, if the target has one.""" - return next( - (level for level in self.explicit_levels if level.name == name), None - ) + return next((level for level in self.explicit_levels if level.name == name), None) def implicit(self, name: str) -> ImplicitMemoryLevelFacts | None: """The implicit level called *name*, if the target has one.""" - return next( - (level for level in self.implicit_levels if level.name == name), None - ) + return next((level for level in self.implicit_levels if level.name == name), None) def cached_level(self, name: str) -> str | None: """The level *name* caches, if it caches one.""" @@ -105,8 +101,7 @@ def cached_level(self, name: str) -> str | None: ( relation.far for relation in self.relations - if relation.kind is MemoryRelationKind.CACHES - and relation.near == name + if relation.kind is MemoryRelationKind.CACHES and relation.near == name ), None, ) @@ -122,9 +117,7 @@ def backing_level(self, name: str) -> str: current = name while (nearer := self.cached_level(current)) is not None: if nearer in seen: - raise ValueError( - f"memory hierarchy: caching cycle through {current!r}" - ) + raise ValueError(f"memory hierarchy: caching cycle through {current!r}") seen.add(nearer) current = nearer return current diff --git a/src/tilefoundry/analysis/memory.py b/src/tilefoundry/analysis/memory.py index 218f8ecd6..4660a9aff 100644 --- a/src/tilefoundry/analysis/memory.py +++ b/src/tilefoundry/analysis/memory.py @@ -277,8 +277,9 @@ def _shares( ( name, Spread( - shares.get(_WHOLE, TrafficBytes()), - tuple(shares.get(unit, TrafficBytes()) for unit in topologies), + logical=shares.get(_WHOLE, TrafficBytes()), + total=shares.get(_WHOLE, TrafficBytes()), + per_unit=tuple(shares.get(unit, TrafficBytes()) for unit in topologies), ), ) for name, shares in sorted(held.items()) diff --git a/src/tilefoundry/analysis/metadata.py b/src/tilefoundry/analysis/metadata.py index ed87f58f2..9e015eecf 100644 --- a/src/tilefoundry/analysis/metadata.py +++ b/src/tilefoundry/analysis/metadata.py @@ -17,60 +17,52 @@ @dataclass(frozen=True) class Spread[V]: - """One quantity, whole and as one unit of each topology level holds it. - - ``per_unit`` runs in the order the record's ``topologies`` states, one - entry per level, so a level's name is written once for the whole record - rather than once per quantity. A finer level's unit sits inside a coarser - one, so its share is the coarser share divided again by whatever the mesh - splits between them -- which is why every level is stated rather than only - the one an analysis was asked about. + """One typed value in logical, expanded, and topology-unit domains. + + ``per_unit`` follows ``ComputeCostMetadata.topologies``. The value retains + its semantic structure while ``Spread`` states only its counting domain. """ + logical: V total: V per_unit: tuple[V, ...] = () def at(self, index: int) -> V: - """One level's share by position, or the total when there is none.""" + """One topology level's value by position, or ``total`` when absent.""" return self.per_unit[index] if index < len(self.per_unit) else self.total @dataclass(frozen=True) class Breakdown[V]: - """One category's quantities, split by the kind of thing each one is. - - The kind is what the rate pricing it is stated for: a dtype prices flops, - a service kind prices what is not floating point, a memory level prices - bytes. Two kinds are never summed, because two rates cannot be. - """ + """Typed quantities grouped by resource kind for traffic records.""" kinds: tuple[tuple[str, Spread[V]], ...] = () def of(self, kind: str) -> Spread[V] | None: - """This kind's quantity, or ``None`` when the record states none.""" return next((value for name, value in self.kinds if name == kind), None) def names(self) -> tuple[str, ...]: - """Every kind this record states, in the order it states them.""" return tuple(name for name, _ in self.kinds) def breakdown[V]( - total: "Mapping[str, V]", per_unit: "Sequence[Mapping[str, V]]", zero: V + total: "Mapping[str, V]", + per_unit: "Sequence[Mapping[str, V]]", + zero: V, + *, + logical: "Mapping[str, V] | None" = None, ) -> Breakdown[V]: - """Gather one category's kinds, each with its total and every level's share. - - A kind any of them states appears in all of them, at *zero* where it was - not stated, so one kind's total and its shares stay one row. - """ - kinds = sorted({*total, *(kind for share in per_unit for kind in share)}) + """Gather keyed spreads, filling missing values in every counting domain.""" + logical = total if logical is None else logical + kinds = sorted({*logical, *total, *(kind for share in per_unit for kind in share)}) return Breakdown( tuple( ( kind, Spread( - total.get(kind, zero), - tuple(share.get(kind, zero) for share in per_unit), + logical=logical.get(kind, zero), + total=total.get(kind, zero), + per_unit=tuple(share.get(kind, zero) for share in per_unit), ), ) for kind in kinds @@ -79,15 +71,11 @@ def breakdown[V]( def shares[V]( - held: Breakdown[V], topologies: tuple[str, ...], topology_level: "str | None" = None + held: Breakdown[V], + topologies: tuple[str, ...], + topology_level: "str | None" = None, ) -> dict[str, V]: - """Each kind's value for one unit of *topology_level*, or its total without one. - - The only place a level's name is turned back into a position, because - ``topologies`` is where the names are written and a ``Spread`` states its - shares in that order and carries none of its own. A level the record does - not state reads as the total, which is what a record over one unit says. - """ + """Each kind's value at one topology level, or its total.""" index = topologies.index(topology_level) if topology_level in topologies else None return { kind: spread.total if index is None else spread.at(index) for kind, spread in held.kinds @@ -96,19 +84,11 @@ def shares[V]( @dataclass(frozen=True) class ComputeCostMetadata(IRMetadata): - """Record one occurrence's work, or one Function's total work. - - ``service`` counts what is not floating point -- comparing, selecting, - whole-number arithmetic -- by the service it asks for, because a predicate - priced as a FLOP is a number about a pipe the work never went down. What an - occurrence moves is a separate record, kept by the family that knows where - values live. On a Call these state one occurrence; on a Function, loops - contribute their trip count. - """ + """Floating-point and other operation kinds, each with one counting spread.""" topologies: tuple[str, ...] = () flops: Breakdown[int] = Breakdown() - service: Breakdown[int] = Breakdown() + other_ops: Breakdown[int] = Breakdown() @dataclass(frozen=True) @@ -129,6 +109,17 @@ class TrafficMetadata(IRMetadata): communication: Breakdown[TrafficBytes] = Breakdown() operands: tuple[TrafficBytes, ...] = () + @property + def whole(self) -> tuple[tuple[str, TrafficBytes], ...]: + return tuple((name, spread.total) for name, spread in self.storage.kinds) + + @property + def per_unit(self) -> tuple[tuple[str, TrafficBytes], ...]: + return tuple( + (name, spread.per_unit[-1] if spread.per_unit else spread.total) + for name, spread in self.storage.kinds + ) + @dataclass(frozen=True) class MemoryLevelFootprint: @@ -143,6 +134,11 @@ class MemoryLevelFootprint: persistent_bytes: int capacity_bytes: int | None = None + @property + def level(self) -> str: + """Compatibility spelling for consumers rendering a memory level.""" + return self.memory_level + @property def exceeds_capacity(self) -> bool: """Whether the peak does not fit the stated capacity.""" @@ -159,6 +155,10 @@ class BufferFootprint: device_bytes: int repeated_bytes: int + @property + def level(self) -> str: + return self.memory_level + @dataclass(frozen=True) class LoopFootprintMetadata(IRMetadata): @@ -191,6 +191,11 @@ class ValueLifetime: last_used_at: int persistent: bool = False + @property + def level(self) -> str: + """Compatibility spelling for consumers rendering a memory level.""" + return self.memory_level + @dataclass(frozen=True) class AllocationMetadata: diff --git a/src/tilefoundry/analysis/report.py b/src/tilefoundry/analysis/report.py index 2e78e83b1..d345da0a6 100644 --- a/src/tilefoundry/analysis/report.py +++ b/src/tilefoundry/analysis/report.py @@ -104,6 +104,7 @@ def _reported_value(value: object, declared: object) -> object: if origin is Spread: (item,) = get_args(declared) return { + "logical": _reported_value(value.logical, item), "total": _reported_value(value.total, item), "per_unit": [_reported_value(share, item) for share in value.per_unit], } @@ -135,6 +136,7 @@ def _pair_types(declared: object) -> tuple[object, object] | None: for _record_type in ( ComputeCostMetadata, + TrafficMetadata, LoopFootprintMetadata, MemoryMetadata, RooflineMetadata, @@ -305,14 +307,12 @@ def _loop_records( return rows -def _totals_of(held: "Breakdown | None") -> dict[str, object]: - """Each kind's whole-program amount, without any level's share. - - A total is what the whole program asks for, so this section states one - number per kind; the per-level shares are on the family's own record. - """ +def _totals_of(held: "Breakdown | Spread | None") -> dict[str, object]: + """Each kind's whole-program amount, without any level's share.""" if held is None: return {} + if isinstance(held, Spread): + return dict(held.total) return { kind: _reported_value(spread.total, type(spread.total)) if is_dataclass(type(spread.total)) diff --git a/src/tilefoundry/analysis/roofline.py b/src/tilefoundry/analysis/roofline.py index 695893c40..6f24a22fd 100644 --- a/src/tilefoundry/analysis/roofline.py +++ b/src/tilefoundry/analysis/roofline.py @@ -31,8 +31,8 @@ def _totals(held) -> tuple[tuple[str, int], ...]: - """One category's kinds with their whole-program counts.""" - return tuple((kind, spread.total) for kind, spread in held.kinds) + """One typed work category's whole-program counts.""" + return tuple((name, spread.total) for name, spread in held.kinds) def _compute_ns(flops: tuple[tuple[str, int], ...], facts: ThroughputFacts) -> int: diff --git a/src/tilefoundry/codegen/context.py b/src/tilefoundry/codegen/context.py index ce521a77e..7f06cd047 100644 --- a/src/tilefoundry/codegen/context.py +++ b/src/tilefoundry/codegen/context.py @@ -12,14 +12,12 @@ from collections.abc import Callable, Mapping from tilefoundry.codegen.signature import CallableSignature, Signature, TensorSignature -from tilefoundry.ir.core.expr import Call from tilefoundry.ir.tir.stmts import Evaluate -from tilefoundry.ir.types import UnitType from tilefoundry.target.base import Target from tilefoundry.visitor_registry.registries import DispatchRegistry, Role, spelled -class CodegenContext: +class EmitContext: """One compile, from the first function read to the last line written.""" target_kind: type[Target] @@ -38,10 +36,14 @@ def __init__( *, symbols: Mapping[int, CallableSignature] | None = None, target: Target | None = None, + resolved_launches: Mapping[int, object] | None = None, + codegen_context: object | None = None, ) -> None: self.registry = registry + self.codegen_context = codegen_context self.symbols: Mapping[int, CallableSignature] = {} if symbols is None else symbols self.target = target + self.resolved_launches = {} if resolved_launches is None else resolved_launches self.exported = False self._lines: list[str] = [] self._indent = 0 @@ -49,6 +51,13 @@ def __init__( self._counter = 0 self._kernel_param_ids: set[int] = set() + def callee_of(self, launch): + """The callee root context resolved for one Launch statement.""" + callee = self.resolved_launches.get(id(launch)) + if callee is None: + raise KeyError("codegen: Launch was not resolved by the root context") + return callee + def signature_of(self, fn) -> CallableSignature: """How a caller writes a call to *fn*, as the symbol table settled it.""" signature = self.symbols.get(id(fn)) @@ -109,14 +118,24 @@ def capture(self, fn) -> str: self._lines = saved_lines def emit_node(self, node) -> None: - """Write *node*, dispatching an ``Evaluate`` by the Op it wraps. + """Write *node* through the target Visitor or function orchestrator.""" + from tilefoundry.ir.tir.abort import Abort # noqa: PLC0415 + from tilefoundry.ir.tir.stmts import ( # noqa: PLC0415 + For, + If, + LetStmt, + MeshScope, + Return, + Sequential, + While, + ) - An effect Op reaches codegen inside an ``Evaluate``, so the handler - that answers for it is registered against the Op's own class. - """ - if isinstance(node, Evaluate): - op = node.callable - self.handler_for(type(op))(Call(type=UnitType(), target=op, args=node.args), self) + if isinstance( + node, (Abort, Evaluate, For, If, LetStmt, MeshScope, Return, Sequential, While) + ): + from tilefoundry.codegen.emitter import CudaEmitter # noqa: PLC0415 + + CudaEmitter(context=self).visit(node) return self.handler_for(type(node))(node, self) @@ -174,9 +193,7 @@ def argument(self, signature: Signature, callee_target: Target) -> tuple[str, .. def arguments(self, signature: CallableSignature, callee_target: Target) -> str: """The whole argument list of one call, in the order the callee declared.""" return ", ".join( - token - for param in signature.all_params - for token in self.argument(param, callee_target) + token for param in signature.all_params for token in self.argument(param, callee_target) ) def local_value(self, signature: Signature) -> str: @@ -204,4 +221,4 @@ def _handler(self, key) -> Callable: return handler -__all__ = ["CodegenContext"] +__all__ = ["EmitContext"] diff --git a/src/tilefoundry/codegen/context_builder.py b/src/tilefoundry/codegen/context_builder.py new file mode 100644 index 000000000..00f827960 --- /dev/null +++ b/src/tilefoundry/codegen/context_builder.py @@ -0,0 +1,312 @@ +"""Build the compile-wide code-generation context before source emission.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, replace +from types import MappingProxyType + +from tilefoundry.codegen.signature import ( + LAUNCH_ABI, + CallableSignature, + called_as, + program_id_params, +) +from tilefoundry.ir.core.module import Module, module_functions, owning_module, subtree +from tilefoundry.ir.tir.launch import Launch +from tilefoundry.ir.tir.prim_function import PrimFunction +from tilefoundry.ir.tir.stmts import Evaluate, Sequential +from tilefoundry.target import CpuTarget, Target + + +@dataclass(frozen=True) +class FunctionSymbol: + """Resolved IR identity and generated ABI names for one callable function.""" + + function: PrimFunction + ir_name: str + kernel_name: str | None + shim_name: str | None + host_name: str | None + owner: Module + target: Target + callable: CallableSignature + + +@dataclass(frozen=True) +class FunctionSymbolTable: + """Separate IR and generated-name namespaces shared by every emitter.""" + + by_function: Mapping[int, FunctionSymbol] + by_ir_name: Mapping[str, FunctionSymbol] + by_kernel_name: Mapping[str, FunctionSymbol] + + @property + def callables(self) -> Mapping[int, CallableSignature]: + return MappingProxyType( + {identity: symbol.callable for identity, symbol in self.by_function.items()} + ) + + +@dataclass(frozen=True) +class LaunchEdge: + """A resolved host-to-device Launch and its domain-owned callable ABI.""" + + launch: Evaluate + caller: PrimFunction + callee: PrimFunction + symbol: FunctionSymbol + callable: CallableSignature + grid_x: object + block_x: object + + +@dataclass(frozen=True) +class EmissionGroup: + owner: Module + target: Target + functions: tuple[PrimFunction, ...] + + +@dataclass(frozen=True) +class EmitContext: + """The compile facts required by one target/topology emission group.""" + + codegen: CodegenContext + owner: Module + target: Target + functions: tuple[PrimFunction, ...] + symbols: Mapping[int, CallableSignature] + launches: Mapping[int, tuple[tuple[object, object, object], tuple[object, object, object]]] + resolved_launches: Mapping[int, PrimFunction] + + +@dataclass(frozen=True) +class CodegenContext: + """Complete compile-wide resolution performed once before emission.""" + + module: Module + functions: tuple[PrimFunction, ...] + ownership: Mapping[int, Module] + targets: Mapping[int, Target] + symbols: FunctionSymbolTable + launch_edges: tuple[LaunchEdge, ...] + topology_domains: Mapping[int, Module] + launch_callables: Mapping[int, CallableSignature] + groups: tuple[EmissionGroup, ...] + + def for_group(self, group: EmissionGroup) -> EmitContext: + identities = {id(fn) for fn in group.functions} + symbols = { + identity: symbol.callable + for identity, symbol in self.symbols.by_function.items() + if identity in identities or symbol.target != group.target + } + launches = { + id(edge.callee): ((edge.grid_x, 1, 1), (edge.block_x, 1, 1)) + for edge in self.launch_edges + if id(edge.callee) in identities + } + for edge in self.launch_edges: + if id(edge.callee) in identities or edge.caller in group.functions: + symbols[id(edge.callee)] = edge.callable + resolved_launches = { + id(edge.launch): edge.callee + for edge in self.launch_edges + if edge.caller in group.functions + } + return EmitContext( + codegen=self, + owner=group.owner, + target=group.target, + functions=group.functions, + symbols=MappingProxyType(symbols), + launches=MappingProxyType(launches), + resolved_launches=MappingProxyType(resolved_launches), + ) + + +def _domain_functions(node: Module) -> tuple[PrimFunction, ...]: + inherited = tuple( + function + for child in node.modules + if child.topologies is None + for function in _domain_functions(child) + ) + return (*node.functions, *inherited) + + +def _domains(root: Module) -> tuple[tuple[Module, tuple[PrimFunction, ...]], ...]: + domains = [(root, _domain_functions(root))] + for node in subtree(root): + domains.extend( + (child, _domain_functions(child)) + for child in node.modules + if child.topologies is not None + ) + return tuple(domains) + + +def _product(values: tuple[object, ...]) -> object: + result: object = 1 + for value in values: + if value is None: + raise ValueError( + "codegen: topology extent cannot be None; use a DimVar-backed " + "CPU-computable launch expression" + ) + result = result * value + return result + + +def _launch_extents(owner: Module) -> tuple[object, object]: + """Derive one-dimensional grid/block extents from declared topologies.""" + ctas = tuple( + topology.size for topology in owner.effective_topologies() if topology.name == "cta" + ) + threads = tuple( + topology.size for topology in owner.effective_topologies() if topology.name == "thread" + ) + return _product(ctas), _product(threads) + + +def _launch_statements(function: PrimFunction) -> tuple[Evaluate, ...]: + body = function.body + if not isinstance(body, Sequential): + return () + return tuple( + stmt + for stmt in body.body + if isinstance(stmt, Evaluate) and isinstance(stmt.callable, Launch) + ) + + +def build_codegen_context(module: Module) -> CodegenContext: + """Resolve functions, symbols, Launch edges, domains, and groups once.""" + functions = tuple(module_functions(module)) + if not all(isinstance(function, PrimFunction) for function in functions): + bad = next(function for function in functions if not isinstance(function, PrimFunction)) + raise TypeError( + f"codegen expects PrimFunction values, got {type(bad).__name__} {bad.name!r}" + ) + expanded_list: list[PrimFunction] = [] + seen_expanded: set[int] = set() + for function in functions: + for variant in function.variants or (function,): + if id(variant) not in seen_expanded: + seen_expanded.add(id(variant)) + expanded_list.append(variant) + expanded = tuple(expanded_list) + ownership: dict[int, Module] = {} + targets: dict[int, Target] = {} + for function in (*functions, *expanded): + owner = owning_module(module, function) + if owner is None: + raise ValueError(f"codegen: function {function.name!r} has no unique owning Module") + if function.target is None: + raise ValueError(f"codegen: function {function.name!r} has no resolved Target") + ownership[id(function)] = owner + targets[id(function)] = function.target + + device_targets = tuple( + dict.fromkeys( + function.target + for function in (*functions, *expanded) + if not isinstance(function.target, CpuTarget) + ) + ) + if len(device_targets) != 1: + raise ValueError(f"codegen: expected one device Target, found {len(device_targets)}") + program_ids = program_id_params(module, device_targets[0]) + by_function: dict[int, FunctionSymbol] = {} + by_ir_name: dict[str, FunctionSymbol] = {} + by_kernel_name: dict[str, FunctionSymbol] = {} + for function in functions: + signature = called_as(function, program_ids) + is_device = not isinstance(function.target, CpuTarget) + symbol = FunctionSymbol( + function=function, + ir_name=function.name, + kernel_name=function.name if is_device else None, + shim_name=signature.name if is_device else None, + host_name=signature.name if not is_device else None, + owner=ownership[id(function)], + target=function.target, + callable=signature, + ) + if function.name in by_ir_name: + raise ValueError(f"codegen: duplicate IR function name {function.name!r}") + by_function[id(function)] = symbol + by_ir_name[function.name] = symbol + if symbol.kernel_name is not None: + if symbol.kernel_name in by_kernel_name: + raise ValueError(f"codegen: duplicate kernel symbol {symbol.kernel_name!r}") + by_kernel_name[symbol.kernel_name] = symbol + + table = FunctionSymbolTable( + MappingProxyType(by_function), + MappingProxyType(by_ir_name), + MappingProxyType(by_kernel_name), + ) + edges: list[LaunchEdge] = [] + launch_callables: dict[int, CallableSignature] = {} + for caller in functions: + for stmt in _launch_statements(caller): + reference = stmt.args[0] + symbol = by_ir_name.get(reference.name) + if symbol is None: + raise ValueError( + f"codegen: launch from {caller.name!r} names unknown function {reference.name!r}" + ) + grid_x, block_x = _launch_extents(symbol.owner) + callable_signature = replace(symbol.callable, trailing=LAUNCH_ABI) + launch_callables[id(stmt)] = callable_signature + edges.append( + LaunchEdge( + stmt, + caller, + symbol.function, + symbol, + callable_signature, + grid_x, + block_x, + ) + ) + + domain_rows = _domains(module) + topology_domains = {id(function): owner for owner, owned in domain_rows for function in owned} + groups: list[EmissionGroup] = [] + for owner, owned in domain_rows: + by_target: dict[Target, list[PrimFunction]] = {} + for function in owned: + by_target.setdefault(function.target, []).append(function) + for target, grouped in by_target.items(): + if not isinstance(target, CpuTarget) and not any( + edge.callee in grouped for edge in edges + ): + missing = ", ".join(function.name for function in grouped) + raise ValueError(f"codegen: no Launch states geometry for {missing}") + groups.append(EmissionGroup(owner, target, tuple(grouped))) + + return CodegenContext( + module=module, + functions=functions, + ownership=MappingProxyType(ownership), + targets=MappingProxyType(targets), + symbols=table, + launch_edges=tuple(edges), + topology_domains=MappingProxyType(topology_domains), + launch_callables=MappingProxyType(launch_callables), + groups=tuple(groups), + ) + + +__all__ = [ + "CodegenContext", + "EmissionGroup", + "EmitContext", + "FunctionSymbol", + "FunctionSymbolTable", + "LaunchEdge", + "build_codegen_context", +] diff --git a/src/tilefoundry/codegen/cpu/abi.py b/src/tilefoundry/codegen/cpu/abi.py index 2b095a0e6..30414d083 100644 --- a/src/tilefoundry/codegen/cpu/abi.py +++ b/src/tilefoundry/codegen/cpu/abi.py @@ -7,7 +7,7 @@ from __future__ import annotations -from tilefoundry.codegen.context import CodegenContext +from tilefoundry.codegen.context import EmitContext from tilefoundry.codegen.signature import ( CallableSignature, ProgramIdSignature, @@ -49,13 +49,13 @@ def _called_as_entry( @register_codegen(CpuTarget, Role.CALLEE, TensorSignature) -def _declare_tensor(sig: TensorSignature, ctx: CodegenContext) -> tuple[str, ...]: +def _declare_tensor(sig: TensorSignature, ctx: EmitContext) -> tuple[str, ...]: """A runtime tensor, which carries its own extents, so nothing else declares them.""" return (f"{TENSOR_CTYPE} {sig.name}",) @register_codegen(CpuTarget, Role.CALLEE, ProgramIdSignature) -def _declare_program_id(sig: ProgramIdSignature, ctx: CodegenContext) -> tuple[str, ...]: +def _declare_program_id(sig: ProgramIdSignature, ctx: EmitContext) -> tuple[str, ...]: """The id of a level no device register answers for, passed in by value.""" return (f"{PROGRAM_ID_CTYPE} {sig.name}",) diff --git a/src/tilefoundry/codegen/cpu/context.py b/src/tilefoundry/codegen/cpu/context.py index 5abe92468..7e9ef7deb 100644 --- a/src/tilefoundry/codegen/cpu/context.py +++ b/src/tilefoundry/codegen/cpu/context.py @@ -10,14 +10,14 @@ from collections.abc import Mapping -from tilefoundry.codegen.context import CodegenContext +from tilefoundry.codegen.context import EmitContext from tilefoundry.codegen.signature import CallableSignature, Signature, TensorSignature from tilefoundry.target import CpuTarget from tilefoundry.target.base import Target from tilefoundry.visitor_registry.registries import codegen_registry -class CpuCodegenContext(CodegenContext): +class CpuCodegenContext(EmitContext): """A compile writing the host unit: the shared context plus what a tensor is here.""" target_kind = CpuTarget @@ -27,8 +27,16 @@ def __init__( *, symbols: Mapping[int, CallableSignature] | None = None, target: Target | None = None, + resolved_launches: Mapping[int, object] | None = None, + codegen_context: object | None = None, ) -> None: - super().__init__(codegen_registry, symbols=symbols, target=target) + super().__init__( + codegen_registry, + symbols=symbols, + target=target, + resolved_launches=resolved_launches, + codegen_context=codegen_context, + ) def local_value(self, signature: Signature) -> str: """The buffer behind a runtime tensor, which is what a device call takes.""" diff --git a/src/tilefoundry/codegen/cpu/module.py b/src/tilefoundry/codegen/cpu/module.py index 54b36a38b..e7d284d46 100644 --- a/src/tilefoundry/codegen/cpu/module.py +++ b/src/tilefoundry/codegen/cpu/module.py @@ -162,9 +162,9 @@ def _as_this_scope_names_it( states keep their names: the entry was told them under exactly those. """ geometry = ( - *(_extent(arg) for arg in launch.args[1:7]), + _extent(launch.args[1]), + _extent(launch.args[4]), str(_static_smem(launch.callable.dynamic_smem)), - "nullptr", ) return replace( shim, @@ -222,10 +222,7 @@ def _declare_callee(callee: PrimFunction, ctx: CpuCodegenContext) -> str: it cannot come to disagree about what the symbol takes. """ shim = ctx.signature_of(callee) - return ( - f'extern "C" void {shim.name}' - f"({ctx.parameters(shim, callee.target, exported=True)});" - ) + return f'extern "C" void {shim.name}({ctx.parameters(shim, callee.target, exported=True)});' def _host_entry( @@ -266,7 +263,7 @@ def emit_host_module( declarations: dict[str, str] = {} body_lines: list[str] = [] for launch in _launch_statements(entry): - callee = module.lookup(launch.args[0].name) + callee = ctx.callee_of(launch) declarations[ctx.signature_of(callee).name] = _declare_callee(callee, ctx) body_lines += _one_launch(entry, launch, callee, ctx) return LinkableModule( diff --git a/src/tilefoundry/codegen/cuda/abi.py b/src/tilefoundry/codegen/cuda/abi.py index ce282307e..cabf6d37c 100644 --- a/src/tilefoundry/codegen/cuda/abi.py +++ b/src/tilefoundry/codegen/cuda/abi.py @@ -8,10 +8,10 @@ from __future__ import annotations -from tilefoundry.codegen.context import CodegenContext +from tilefoundry.codegen.context import EmitContext from tilefoundry.codegen.cuda.context import CudaCodegenContext +from tilefoundry.codegen.cuda.tir.stmts.mesh_scope import mesh_type from tilefoundry.codegen.signature import ( - LAUNCH_ABI, CallableSignature, LaunchSignature, ProgramIdSignature, @@ -21,6 +21,7 @@ tensor_signature_of, ) from tilefoundry.ir.tir.prim_function import PrimFunction +from tilefoundry.ir.types.shard.shard_layout import ShardLayout from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.registries import Role, register_codegen @@ -85,9 +86,10 @@ def kernel_arguments(kernel: CallableSignature, ctx: CudaCodegenContext) -> str: continue element = ctx.dtype_to_cpp(param.type.dtype.name) written.append(f"static_cast<{element}*>({param.name})") - written += [ - f"static_cast({param.extent_name(axis)})" for axis in param.dynamic_axes - ] + written += [f"static_cast({param.extent_name(axis)})" for axis in param.dynamic_axes] + layout = param.type.layout + if isinstance(layout, ShardLayout): + written.append(f"{mesh_type(layout.mesh)}{{}}") return ", ".join(written) @@ -106,12 +108,11 @@ def _called_as_shim( params=tuple(tensor_signature_of(var) for var in fn.params), output_count=fn.output_count, leading=program_ids, - trailing=LAUNCH_ABI, ) @register_codegen(CudaTarget, Role.CALLEE, TensorSignature) -def _declare_tensor(sig: TensorSignature, ctx: CodegenContext) -> tuple[str, ...]: +def _declare_tensor(sig: TensorSignature, ctx: EmitContext) -> tuple[str, ...]: """A tensor arrives as a pointer, then one extent per open axis. A static extent is a constant on both sides of the call and costs the @@ -123,10 +124,14 @@ def _declare_tensor(sig: TensorSignature, ctx: CodegenContext) -> tuple[str, ... pointer, extent = _C_ABI else: pointer, extent = f"{ctx.dtype_to_cpp(sig.type.dtype.name)}*", "int" - return ( + declared = [ f"{pointer} {sig.name}", *(f"{extent} {sig.extent_name(axis)}" for axis in sig.dynamic_axes), - ) + ] + layout = sig.type.layout + if not ctx.exported and isinstance(layout, ShardLayout): + declared.append(f"{mesh_type(layout.mesh)} {sig.name}_mesh") + return tuple(declared) @register_codegen(CudaTarget, Role.CALLEE, LaunchSignature) @@ -148,7 +153,7 @@ def _declare_program_meta(sig: ProgramMetaSignature, ctx: CudaCodegenContext) -> @register_codegen(CudaTarget, Role.CALLER, TensorSignature) -def _pass_tensor(sig: TensorSignature, ctx: CodegenContext) -> tuple[str, ...]: +def _pass_tensor(sig: TensorSignature, ctx: EmitContext) -> tuple[str, ...]: """The pointer, then the extents, as the calling scope spells them. How many arguments there are and in what order is CUDA's, since CUDA @@ -163,7 +168,7 @@ def _pass_tensor(sig: TensorSignature, ctx: CodegenContext) -> tuple[str, ...]: @register_codegen(CudaTarget, Role.CALLER, LaunchSignature) @register_codegen(CudaTarget, Role.CALLER, ProgramIdSignature) @register_codegen(CudaTarget, Role.CALLER, ProgramMetaSignature) -def _pass_value(sig: Signature, ctx: CodegenContext) -> tuple[str, ...]: +def _pass_value(sig: Signature, ctx: EmitContext) -> tuple[str, ...]: """One argument the caller already holds, spelled in its own scope.""" return (ctx.local_value(sig),) diff --git a/src/tilefoundry/codegen/cuda/context.py b/src/tilefoundry/codegen/cuda/context.py index fe4e673d9..7ad734de5 100644 --- a/src/tilefoundry/codegen/cuda/context.py +++ b/src/tilefoundry/codegen/cuda/context.py @@ -11,14 +11,15 @@ from collections.abc import Iterable, Mapping -from tilefoundry.codegen.context import CodegenContext +from tilefoundry.codegen.context import EmitContext from tilefoundry.codegen.signature import CallableSignature, TensorSignature -from tilefoundry.codegen.topology import Geometry from tilefoundry.ir.types.dim import DimVar from tilefoundry.target import CudaTarget from tilefoundry.target.base import Target from tilefoundry.visitor_registry.registries import codegen_registry +Geometry = tuple[tuple[object, object, object], tuple[object, object, object]] + _CUDA_CPP: dict[str, str] = { "f32": "float", "bf16": "__nv_bfloat16", @@ -47,7 +48,7 @@ def topology_scope_str(name: str) -> str: ) from None -class CudaCodegenContext(CodegenContext): +class CudaCodegenContext(EmitContext): """A compile writing CUDA: the shared context plus what only CUDA states.""" target_kind = CudaTarget @@ -58,8 +59,14 @@ def __init__( symbols: Mapping[int, CallableSignature] | None = None, target: Target | None = None, launches: Mapping[int, Geometry] | None = None, + codegen_context: object | None = None, ) -> None: - super().__init__(codegen_registry, symbols=symbols, target=target) + super().__init__( + codegen_registry, + symbols=symbols, + target=target, + codegen_context=codegen_context, + ) self._mesh_aliases: dict[int, tuple[str, str]] = {} self.launches: Mapping[int, Geometry] = {} if launches is None else launches """The geometry each device function is launched at, keyed by ``id(fn)``.""" diff --git a/src/tilefoundry/codegen/cuda/module.py b/src/tilefoundry/codegen/cuda/module.py index 9be4c143b..a01bb7c19 100644 --- a/src/tilefoundry/codegen/cuda/module.py +++ b/src/tilefoundry/codegen/cuda/module.py @@ -23,11 +23,13 @@ from tilefoundry.codegen.linkable import LinkableFunction, LinkableModule from tilefoundry.codegen.registry import CodeGenerator from tilefoundry.codegen.signature import CallableSignature -from tilefoundry.codegen.topology import Geometry, coarsest_topology from tilefoundry.ir.core.module import Module from tilefoundry.ir.tir.prim_function import PrimFunction +from tilefoundry.ir.types.shape_helpers import static_dim_value from tilefoundry.target import Target +Geometry = tuple[tuple[object, object, object], tuple[object, object, object]] + def _emit_function(fn: PrimFunction, ctx: CudaCodegenContext) -> LinkableFunction: """*fn* in both of its positions, written from the one convention it is called by. @@ -56,9 +58,7 @@ def _define( _define_shim(shim, kernel, ctx) -def _define_kernel( - fn: PrimFunction, kernel: CallableSignature, ctx: CudaCodegenContext -) -> None: +def _define_kernel(fn: PrimFunction, kernel: CallableSignature, ctx: CudaCodegenContext) -> None: """The ``__global__``: what it was told, then what it was written to do.""" ctx.reset_barrier_ids() ctx.emit(f"__global__ void {kernel.name}({ctx.parameters(kernel)}) {{") @@ -77,15 +77,10 @@ def _define_shim( ctx.emit(f'extern "C" void {shim.name}({ctx.parameters(shim, exported=True)}) {{') ctx.indent() _tell_the_block(shim, ctx) - grid_x, grid_y, grid_z, block_x, block_y, block_z, smem, stream = ( - signature.name for signature in shim.trailing - ) - ctx.emit(f"dim3 grid({grid_x}, {grid_y}, {grid_z});") - ctx.emit(f"dim3 block({block_x}, {block_y}, {block_z});") - ctx.emit( - f"{kernel.name}<<>>" - f"({kernel_arguments(kernel, ctx)});" - ) + grid_x, block_x, smem = (signature.name for signature in shim.trailing) + ctx.emit(f"dim3 grid({grid_x}, 1, 1);") + ctx.emit(f"dim3 block({block_x}, 1, 1);") + ctx.emit(f"{kernel.name}<<>>({kernel_arguments(kernel, ctx)});") ctx.dedent() ctx.emit("}") @@ -102,8 +97,7 @@ def _tell_the_block(shim: CallableSignature, ctx: CudaCodegenContext) -> None: for signature in shim.leading: scope = topology_scope_str(signature.topology_level) ctx.emit( - f"{PROGRAM_META.name}.program_id[int({scope})] = " - f"static_cast({signature.name});" + f"{PROGRAM_META.name}.program_id[int({scope})] = static_cast({signature.name});" ) @@ -143,9 +137,15 @@ def _program_dims(geometry: Geometry) -> list[dict[str, str]]: """ grid, block = geometry dims = [] - if grid[0] is not None: - dims.append({"scope": topology_scope_str("cta"), "count": str(math.prod(grid))}) - dims.append({"scope": topology_scope_str("thread"), "count": str(math.prod(block))}) + static_grid = tuple(static_dim_value(value) for value in grid) + if all(value is not None for value in static_grid): + dims.append({"scope": topology_scope_str("cta"), "count": str(math.prod(static_grid))}) + elif static_grid[1:] != (1, 1): + raise ValueError("codegen: only grid.x may be a host-computed dynamic expression") + static_block = tuple(static_dim_value(value) for value in block) + if any(value is None for value in static_block): + raise ValueError("codegen: thread topology extents must be compile-time constants") + dims.append({"scope": topology_scope_str("thread"), "count": str(math.prod(static_block))}) return dims @@ -160,14 +160,14 @@ def emit_cuda_module( geometry = _one_geometry(kernels, ctx) preamble = render( "cuda_preamble.cu.j2", - program_topology=topology_scope_str(coarsest_topology(module)), + program_topology=topology_scope_str( + module.effective_topologies()[0].name if module.effective_topologies() else "cta" + ), program_dims=_program_dims(geometry), - dynamic_cta=geometry[0][0] is None, + dynamic_cta=static_dim_value(geometry[0][0]) is None, grid_barrier_state=ctx.needs_grid_barrier_state, ) - return LinkableModule( - target="cuda", language="cu", preamble=preamble, functions=functions - ) + return LinkableModule(target="cuda", language="cu", preamble=preamble, functions=functions) CUDA_CODE_GENERATOR = CodeGenerator(emit_cuda_module) diff --git a/src/tilefoundry/codegen/cuda/tir/prim_function.py b/src/tilefoundry/codegen/cuda/tir/prim_function.py index c3124c312..8ee019cc6 100644 --- a/src/tilefoundry/codegen/cuda/tir/prim_function.py +++ b/src/tilefoundry/codegen/cuda/tir/prim_function.py @@ -11,6 +11,7 @@ from tilefoundry.codegen.cuda.context import CudaCodegenContext from tilefoundry.codegen.cuda.tir.memory.tensor_view import render_shard_layout_value +from tilefoundry.codegen.emitter import CudaEmitter from tilefoundry.codegen.signature import TensorSignature, tensor_signature_of from tilefoundry.ir.core.pattern import DimVarRangePat from tilefoundry.ir.tir.prim_function import PrimFunction @@ -40,10 +41,7 @@ def _wrap_buffer(signature: TensorSignature, ctx: CudaCodegenContext) -> None: ) for line in preamble: ctx.emit(line) - ctx.emit( - f"auto {name}_tensor = tilefoundry::make_shard_tensor(" - f"{pointer}, {layout}, {value});" - ) + ctx.emit(f"auto {name}_tensor = tilefoundry::make_shard_tensor({pointer}, {layout}, {value});") def _dispatch(fn: PrimFunction, ctx: CudaCodegenContext) -> None: @@ -104,4 +102,4 @@ def _emit(fn: PrimFunction, ctx: CudaCodegenContext) -> None: ctx.bind_extents(params) for signature in params: _wrap_buffer(signature, ctx) - ctx.emit_node(fn.body) + CudaEmitter(context=ctx).visit(fn.body) diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/abort.py b/src/tilefoundry/codegen/cuda/tir/stmts/abort.py deleted file mode 100644 index e0e2d42f2..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/abort.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Emit the TIR abort effect.""" - -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.tir.abort import Abort -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - - -@register_codegen(CudaTarget, Role.EMIT, Abort) -def _emit(node: Abort, ctx: CudaCodegenContext) -> None: - ctx.emit("__trap();") diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/for_.py b/src/tilefoundry/codegen/cuda/tir/stmts/for_.py deleted file mode 100644 index 56967de56..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/for_.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Emitter for ``tir.For`` — C-style for loop.""" - -from __future__ import annotations - -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.tir.stmts import For -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - -from .scalar_expr import render_scalar_expr - - -@register_codegen(CudaTarget, Role.EMIT, For) -def _emit(node: For, ctx: CudaCodegenContext) -> None: - iv_name = ctx.name_for(node.induction_var) - start = render_scalar_expr(node.start, ctx) - stop = render_scalar_expr(node.stop, ctx) - step = render_scalar_expr(node.step, ctx) - - ctx.emit(f"for (int {iv_name} = {start}; {iv_name} < {stop}; {iv_name} += {step}) {{") - ctx.indent() - ctx.emit_node(node.body) - ctx.dedent() - ctx.emit("}") diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/if_.py b/src/tilefoundry/codegen/cuda/tir/stmts/if_.py deleted file mode 100644 index 484eed0b5..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/if_.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Emit TIR conditionals with scalar C-style predicates.""" - -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.tir.stmts import If -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - -from .scalar_expr import render_scalar_expr - - -@register_codegen(CudaTarget, Role.EMIT, If) -def _emit(node: If, ctx: CudaCodegenContext) -> None: - cond = render_scalar_expr(node.cond, ctx) - ctx.emit(f"if ({cond}) {{") - ctx.indent() - ctx.emit_node(node.then_body) - ctx.dedent() - if getattr(node.else_body, "body", None): - ctx.emit("} else {") - ctx.indent() - ctx.emit_node(node.else_body) - ctx.dedent() - ctx.emit("}") - else: - ctx.emit("}") diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/let_stmt.py b/src/tilefoundry/codegen/cuda/tir/stmts/let_stmt.py deleted file mode 100644 index 8ed9e2bae..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/let_stmt.py +++ /dev/null @@ -1,30 +0,0 @@ -"""Emitter for ``tir.LetStmt``. - -LetStmt wraps a value binding: `auto %var = ` then the scoped -body. Dispatch is two-step: - -1. Look up an emitter keyed on the ``Op`` class of ``value.target`` — each - TIR-owned Expr Op (``tir.memory.AllocTensor``, ``tir.memory.PtrOf``, - ``tir.memory.TensorView``, etc.) registers its own emitter with - signature ``(let: LetStmt, ctx) -> None``. -2. Recurse into the Sequential body. -""" -from __future__ import annotations - -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.core import Call -from tilefoundry.ir.tir.stmts import LetStmt -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - - -@register_codegen(CudaTarget, Role.EMIT, LetStmt) -def _emit(node: LetStmt, ctx: CudaCodegenContext) -> None: - if not isinstance(node.value, Call): - raise RuntimeError( - f"LetStmt.value must be a Call (TIR-owned Expr Op), " - f"got {type(node.value).__name__}" - ) - op = node.value.target - ctx.handler_for(type(op))(node, ctx) - ctx.emit_node(node.body) diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py b/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py index 0e6ef022c..d7f178633 100644 --- a/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py +++ b/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py @@ -1,21 +1,15 @@ -"""Emit CUDA mesh scopes. +"""Helpers for CUDA ``MeshScope`` lowering. -Emitter for `tir.MeshScope` — emits a C++ block + comment marker + -constexpr Mesh type alias ([runtime §2.3.2](docs/spec/runtime.md#232-layoutmeshcuh)). +The ``CudaEmitter`` owns Stmt traversal and invokes these stateless mesh +layout helpers while visiting a ``MeshScope``. """ - from __future__ import annotations -from tilefoundry.codegen.cuda.context import ( - CudaCodegenContext, - topology_scope_str, -) -from tilefoundry.ir.tir.stmts import MeshScope +from tilefoundry.codegen.cuda.context import topology_scope_str from tilefoundry.ir.tir.sync import participation from tilefoundry.ir.types.shard.layout import ComposedLayout, Layout from tilefoundry.ir.types.shard.mesh import Mesh, Topology, positions_at -from tilefoundry.target import CudaTarget, validate_cuda_topology_levels -from tilefoundry.visitor_registry.registries import Role, register_codegen +from tilefoundry.target import validate_cuda_topology_levels def _resolved(topology: Topology | str) -> Topology: @@ -43,14 +37,14 @@ def _validate_topology(mesh: Mesh, target) -> None: validate_cuda_topology_levels(target, (_resolved(t).name for t in mesh.topologies)) -def _group_mesh_layout(mesh: Mesh, topos) -> str: +def _levelwise_layout(mesh: Mesh, topos) -> str: """One nest per topology level, each in that level's own numbering. A flat shape says nothing about which axes are whose, so the runtime's - ``get`` needs the boundary stated. ``topology_axes`` hands the - axes to the levels left to right and ``positions_at`` divides each level's - strides by what the levels under it contribute, so a nest reads as the - layout that level would have alone. The composite is a grouping, not a map: only the + ``get`` needs the boundary stated. ``level_axes`` hands the axes to + the levels left to right and ``positions_at`` divides each level's strides + by what the levels under it contribute, so a nest reads as the layout that + level would have alone. The composite is a grouping, not a map: only the nests are evaluated. A level every instance shares keeps a mode of one, so that ``get`` still has something to pick. """ @@ -59,9 +53,20 @@ def _group_mesh_layout(mesh: Mesh, topos) -> str: level_shape, level_strides = positions_at(mesh, topology.name) if not level_shape: level_shape, level_strides = (1,), (0,) - shapes.append("cute::Shape<" + ", ".join(f"cute::Int<{s}>" for s in level_shape) + ">") - strides.append("cute::Stride<" + ", ".join(f"cute::Int<{s}>" for s in level_strides) + ">") - return f"cute::Layout, cute::Stride<{', '.join(strides)}>>" + shapes.append( + "cute::Shape<" + + ", ".join(f"cute::Int<{s}>" for s in level_shape) + + ">" + ) + strides.append( + "cute::Stride<" + + ", ".join(f"cute::Int<{s}>" for s in level_strides) + + ">" + ) + return ( + f"cute::Layout, " + f"cute::Stride<{', '.join(strides)}>>" + ) def mesh_type(mesh: Mesh) -> str: @@ -83,7 +88,10 @@ def mesh_type(mesh: Mesh) -> str: if len(topos) == 1: shape_types = ", ".join(f"cute::Int<{s}>" for s in shape) stride_types = ", ".join(f"cute::Int<{s}>" for s in strides) - layout = f"cute::Layout, cute::Stride<{stride_types}>>" + layout = ( + f"cute::Layout, " + f"cute::Stride<{stride_types}>>" + ) else: if base: raise NotImplementedError( @@ -91,7 +99,7 @@ def mesh_type(mesh: Mesh) -> str: "be sliced; the slice and the level boundary would both be " "deciding which positions these are" ) - layout = _group_mesh_layout(mesh, topos) + layout = _levelwise_layout(mesh, topos) if base: layout = f"cute::ComposedLayout, {layout}>" return f"tilefoundry::Mesh<{layout}, {', '.join(topology_scope_str(t.name) for t in topos)}>" @@ -108,35 +116,4 @@ def _is_dynamic_mesh(mesh: Mesh) -> bool: return any(s is None for s in mesh.layout.shape) -@register_codegen(CudaTarget, Role.EMIT, MeshScope) -def _emit(node: MeshScope, ctx: CudaCodegenContext) -> None: - """Emit the block a mesh scope is, and the mesh object it states.""" - if ctx.target is None: - raise RuntimeError("CUDA MeshScope emission requires its Target") - _validate_topology(node.mesh, ctx.target) - name = ctx.name_for(node.binding) - ctx.emit(f"// mesh scope: {program_topologies(node.mesh)[0].name}") - - is_slice = isinstance(node.mesh.layout, ComposedLayout) - ctx.emit("{") - ctx.indent() - outer_aliases = ctx._mesh_aliases - ctx._mesh_aliases = dict(outer_aliases) - try: - if not _is_dynamic_mesh(node.mesh): - alias = f"{name}_mesh_t" - mesh_type_str = mesh_type(node.mesh) - ctx._mesh_aliases[id(node.mesh)] = (alias, mesh_type_str) - ctx.emit(f"using {alias} = {mesh_type_str};") - ctx.emit(f"constexpr {alias} {name}_mesh{{}};") - if is_slice: - ctx.emit(f"if (tilefoundry::contains({name}_mesh, tilefoundry::program_ids())) {{") - ctx.indent() - ctx.emit_node(node.body) - if is_slice: - ctx.dedent() - ctx.emit("}") - finally: - ctx._mesh_aliases = outer_aliases - ctx.dedent() - ctx.emit("}") +__all__ = ["mesh_type", "program_topologies"] diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/return_.py b/src/tilefoundry/codegen/cuda/tir/stmts/return_.py deleted file mode 100644 index ea05e6271..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/return_.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Emitter for `tir.Return`.""" - -from __future__ import annotations - -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.tir.stmts import Return -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - - -@register_codegen(CudaTarget, Role.EMIT, Return) -def _emit(node: Return, ctx: CudaCodegenContext) -> None: - ctx.emit("return;") diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/sequential.py b/src/tilefoundry/codegen/cuda/tir/stmts/sequential.py deleted file mode 100644 index fadc59595..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/sequential.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Emitter for ``tir.Sequential`` — packs a list of Stmts in order.""" -from __future__ import annotations - -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.tir.stmts import Sequential -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - - -@register_codegen(CudaTarget, Role.EMIT, Sequential) -def _emit(node: Sequential, ctx: CudaCodegenContext) -> None: - for stmt in node.body: - ctx.emit_node(stmt) diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/while_.py b/src/tilefoundry/codegen/cuda/tir/stmts/while_.py deleted file mode 100644 index df673bc99..000000000 --- a/src/tilefoundry/codegen/cuda/tir/stmts/while_.py +++ /dev/null @@ -1,15 +0,0 @@ -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.ir.tir.stmts import While -from tilefoundry.target import CudaTarget -from tilefoundry.visitor_registry.registries import Role, register_codegen - -from .scalar_expr import render_scalar_expr - - -@register_codegen(CudaTarget, Role.EMIT, While) -def _emit(node: While, ctx: CudaCodegenContext) -> None: - ctx.emit(f"while ({render_scalar_expr(node.cond, ctx)}) {{") - ctx.indent() - ctx.emit_node(node.body) - ctx.dedent() - ctx.emit("}") diff --git a/src/tilefoundry/codegen/emitter.py b/src/tilefoundry/codegen/emitter.py new file mode 100644 index 000000000..dcd5ce894 --- /dev/null +++ b/src/tilefoundry/codegen/emitter.py @@ -0,0 +1,158 @@ +"""C++ source emitters shared by CUDA and CPU code generation.""" + +from __future__ import annotations + +from tilefoundry.codegen.cuda.tir.stmts.mesh_scope import ( + _is_dynamic_mesh, + _validate_topology, + mesh_type, + program_topologies, +) +from tilefoundry.codegen.cuda.tir.stmts.scalar_expr import render_scalar_expr +from tilefoundry.ir.core import Call +from tilefoundry.ir.tir.abort import Abort +from tilefoundry.ir.tir.stmts import ( + Evaluate, + For, + If, + LetStmt, + MeshScope, + Return, + Sequential, + While, +) +from tilefoundry.ir.types import UnitType +from tilefoundry.ir.types.shard.layout import ComposedLayout +from tilefoundry.ir.visitor import StmtVisitor + + +class CppEmitter(StmtVisitor[None]): + """Write C++ from TIR; the context carries state, this class traverses.""" + + def __init__(self, *, context) -> None: + self.context = context + + def emit(self, line: str) -> None: + self.context.emit(line) + + def indent(self) -> None: + self.context.indent() + + def dedent(self) -> None: + self.context.dedent() + + def statement(self, text: str) -> None: + self.emit(f"{text};") + + def generic_visit(self, node) -> None: + raise RuntimeError(f"no emitter for {type(node).__name__}") + + def visit_Sequential(self, node: Sequential) -> None: + for stmt in node.body: + self.visit(stmt) + + def visit_If(self, node: If) -> None: + cond = self.render_scalar(node.cond) + self.emit(f"if ({cond}) {{") + self.indent() + self.visit(node.then_body) + self.dedent() + if getattr(node.else_body, "body", None): + self.emit("} else {") + self.indent() + self.visit(node.else_body) + self.dedent() + self.emit("}") + + def visit_For(self, node: For) -> None: + name = self.context.name_for(node.induction_var) + start = self.render_scalar(node.start) + stop = self.render_scalar(node.stop) + step = self.render_scalar(node.step) + self.emit(f"for (int {name} = {start}; {name} < {stop}; {name} += {step}) {{") + self.indent() + self.visit(node.body) + self.dedent() + self.emit("}") + + def visit_While(self, node: While) -> None: + self.emit(f"while ({self.render_scalar(node.cond)}) {{") + self.indent() + self.visit(node.body) + self.dedent() + self.emit("}") + + def visit_Return(self, node: Return) -> None: + self.statement("return") + + def visit_Abort(self, node: Abort) -> None: + self.statement("__trap()") + + def visit_LetStmt(self, node: LetStmt) -> None: + if not isinstance(node.value, Call): + raise RuntimeError( + f"LetStmt.value must be a Call (TIR-owned Expr Op), got {type(node.value).__name__}" + ) + self.emit_let_op(node.value, node) + self.visit(node.body) + + def visit_Evaluate(self, node: Evaluate) -> None: + op = node.callable + self.emit_evaluate_op(Call(type=UnitType(), target=op, args=node.args)) + + def visit_MeshScope(self, node: MeshScope) -> None: + raise NotImplementedError + + def emit_evaluate_op(self, call: Call) -> None: + handler = self.context.handler_for(type(call.target)) + if handler is None: + raise RuntimeError(f"no codegen handler for Op {type(call.target).__name__}") + handler(call, self.context) + + def emit_let_op(self, call: Call, stmt: LetStmt) -> None: + handler = self.context.handler_for(type(call.target)) + if handler is None: + raise RuntimeError(f"no codegen handler for Op {type(call.target).__name__}") + handler(stmt, self.context) + + def render_scalar(self, expr) -> str: + return render_scalar_expr(expr, self.context) + + +class CudaEmitter(CppEmitter): + """CUDA C++ emitter, including mesh scope lowering.""" + + def visit_MeshScope(self, node: MeshScope) -> None: + ctx = self.context + if ctx.target is None: + raise RuntimeError("CUDA MeshScope emission requires its Target") + _validate_topology(node.mesh, ctx.target) + name = ctx.name_for(node.binding) + ctx.emit(f"// mesh scope: {program_topologies(node.mesh)[0].name}") + + is_slice = isinstance(node.mesh.layout, ComposedLayout) + ctx.emit("{") + ctx.indent() + outer_aliases = ctx._mesh_aliases + ctx._mesh_aliases = dict(outer_aliases) + try: + if not _is_dynamic_mesh(node.mesh): + alias = f"{name}_mesh_t" + mesh_type_str = mesh_type(node.mesh) + ctx._mesh_aliases[id(node.mesh)] = (alias, mesh_type_str) + ctx.emit(f"using {alias} = {mesh_type_str};") + ctx.emit(f"constexpr {alias} {name}_mesh{{}};") + if is_slice: + ctx.emit(f"if (tilefoundry::contains({name}_mesh, tilefoundry::program_ids())) {{") + ctx.indent() + self.visit(node.body) + if is_slice: + ctx.dedent() + ctx.emit("}") + finally: + ctx._mesh_aliases = outer_aliases + ctx.dedent() + ctx.emit("}") + + +__all__ = ["CppEmitter", "CudaEmitter"] diff --git a/src/tilefoundry/codegen/registry.py b/src/tilefoundry/codegen/registry.py index 5b3048611..3f95cd0d4 100644 --- a/src/tilefoundry/codegen/registry.py +++ b/src/tilefoundry/codegen/registry.py @@ -2,53 +2,6 @@ from __future__ import annotations -from tilefoundry.ir.core.module import Module, module_functions -from tilefoundry.target.base import Target, _target_summary from tilefoundry.target.services import CodeGenerator - -def group_functions_by_target( - module: Module, -) -> dict[Target, tuple["PrimFunction", ...]]: - """Group functions by their equal Target value, preserving source order.""" - from tilefoundry.ir.tir.prim_function import PrimFunction # noqa: PLC0415 - - groups: dict[Target, list[PrimFunction]] = {} - for function in module_functions(module): - if not isinstance(function, PrimFunction): - raise TypeError( - f"tilefoundry: codegen expects PrimFunction values, got " - f"{type(function).__name__} {function.name!r}" - ) - if function.target is None: - raise ValueError( - f"tilefoundry: function {function.name!r} has no resolved Target " - "at codegen grouping" - ) - groups.setdefault(function.target, []).append(function) - """Expose structural variants to their target-specific emitter.""" - for variant in function.variants: - if variant.target is None: - raise ValueError(f"tilefoundry: variant {variant.name!r} has no resolved Target") - groups.setdefault(variant.target, []).append(variant) - - from tilefoundry.target import CudaTarget # noqa: PLC0415 - - device_groups = [ - (target, functions) - for target, functions in groups.items() - if isinstance(target, CudaTarget) - ] - if len(device_groups) > 1: - first_target, first_functions = device_groups[0] - second_target, second_functions = device_groups[1] - raise ValueError( - f"tilefoundry: module {module.name!r} mixes unequal device Targets: " - f"{_target_summary(first_target)} (function {first_functions[0].name!r}) " - f"vs {_target_summary(second_target)} (function {second_functions[0].name!r}); " - "multiple device translation units are not supported" - ) - return {target: tuple(functions) for target, functions in groups.items()} - - -__all__ = ["CodeGenerator", "group_functions_by_target"] +__all__ = ["CodeGenerator"] diff --git a/src/tilefoundry/codegen/signature.py b/src/tilefoundry/codegen/signature.py index 64266be23..bf28c66a1 100644 --- a/src/tilefoundry/codegen/signature.py +++ b/src/tilefoundry/codegen/signature.py @@ -9,10 +9,9 @@ from __future__ import annotations -from collections.abc import Mapping from dataclasses import dataclass -from tilefoundry.ir.core.module import Module, module_functions +from tilefoundry.ir.core.module import Module from tilefoundry.ir.tir.prim_function import PrimFunction from tilefoundry.ir.types import TensorType from tilefoundry.ir.types.dim import DimVar @@ -128,15 +127,14 @@ def all_params(self) -> tuple[Signature, ...]: LAUNCH_ABI = ( LaunchSignature("grid_x", "int"), - LaunchSignature("grid_y", "int"), - LaunchSignature("grid_z", "int"), LaunchSignature("block_x", "int"), - LaunchSignature("block_y", "int"), - LaunchSignature("block_z", "int"), LaunchSignature("dynamic_smem", "int"), - LaunchSignature("stream", "void*"), ) -"""What a launch shim is told beyond the kernel's own arguments.""" +"""Actual cross-translation-unit scalars on a CPU-to-device Launch edge. + +The topology is one-dimensional, so only grid_x and block_x are stated. +Dynamic shared memory remains because its per-Launch value crosses the shim ABI. +""" def tensor_signature_of(var) -> TensorSignature: @@ -187,18 +185,6 @@ def called_as(fn, program_ids: tuple[ProgramIdSignature, ...] = ()) -> CallableS return convention(fn, program_ids) -def symbol_table(module: Module, target: Target) -> Mapping[int, CallableSignature]: - """What every function in the tree is called by, before anything is emitted. - - Keyed by ``id(fn)``: a call site holds the callee itself and never asks - which module it lives in. A specialization variant gets no row, because - nothing calls one: its prototype compiles to the one symbol, and which - variant runs is decided inside it. - """ - program_ids = program_id_params(module, target) - return {id(fn): called_as(fn, program_ids) for fn in module_functions(module)} - - __all__ = [ "LAUNCH_ABI", "CallableSignature", @@ -211,6 +197,5 @@ def symbol_table(module: Module, target: Target) -> Mapping[int, CallableSignatu "UnitSignature", "called_as", "program_id_params", - "symbol_table", "tensor_signature_of", ] diff --git a/src/tilefoundry/codegen/topology.py b/src/tilefoundry/codegen/topology.py deleted file mode 100644 index adc5df63d..000000000 --- a/src/tilefoundry/codegen/topology.py +++ /dev/null @@ -1,112 +0,0 @@ -"""The topology questions every emitter asks of a module and its target. - -Host module, device module and the entry signature describe three segments of -one call, so they read the answer here rather than each deciding for -themselves. -""" - -from __future__ import annotations - -from collections.abc import Iterator, Mapping - -from tilefoundry.ir.core.module import Module, module_functions, subtree -from tilefoundry.ir.tir.launch import Launch -from tilefoundry.ir.tir.stmts import Evaluate, Sequential -from tilefoundry.ir.types.shape_helpers import static_dim_value - -Geometry = tuple[tuple[int | None, int, int], tuple[int, int, int]] -"""The grid and block one launch runs at; a launch-provided CTA extent is ``None``.""" - - -def coarsest_topology(module) -> str: - """The coarsest topology level this module's program names. - - ``effective_topologies()`` is ordered coarsest first, so the first entry - says where this program's run of levels begins; a module that names none - is one CTA. - """ - declared = module.effective_topologies() - return declared[0].name if declared else "cta" - - -def topology_domains(root: Module) -> Iterator[tuple[Module, tuple]]: - """Each module that states topologies of its own, and the functions under it. - - A program states its instance counts once, so a module that states its own - is its own translation unit; one that states none runs in its owner's - program and its functions are emitted into the owner's unit. - """ - yield root, _domain_functions(root) - for node in subtree(root): - for child in node.modules: - if child.topologies is not None: - yield child, _domain_functions(child) - - -def _domain_functions(node: Module) -> tuple: - """*node*'s own functions and those of every child that states no topologies.""" - inherited = tuple( - function - for child in node.modules - if child.topologies is None - for function in _domain_functions(child) - ) - return (*node.functions, *inherited) - - -def launch_geometry(root: Module) -> Mapping[int, Geometry]: - """The geometry each launched function runs at, as its ``Launch`` states it. - - The host entry settled it when it was written, so the device side reads it - rather than walking a body to derive the same answer a second time. - """ - geometry: dict[int, Geometry] = {} - for function in module_functions(root): - body = function.body - if not isinstance(body, Sequential): - continue - for stmt in body.body: - if not (isinstance(stmt, Evaluate) and isinstance(stmt.callable, Launch)): - continue - callee = root.lookup(stmt.args[0].name) - geometry[id(callee)] = ( - _grid(stmt.args[1:4], callee.name), - _block(stmt.args[4:7], callee.name), - ) - return geometry - - -def _grid(args, name: str) -> tuple: - """The three grid extents, whose first may be the one a launch provides. - - A launch-provided CTA count is the one extent a translation unit cannot - state; the runtime states that one when it is asked. - """ - return ( - static_dim_value(args[0]), - *(_static(arg, name, "grid", axis + 1) for axis, arg in enumerate(args[1:])), - ) - - -def _block(args, name: str) -> tuple: - """The three block extents, which multiply into the thread count a unit states.""" - return tuple(_static(arg, name, "block", axis) for axis, arg in enumerate(args)) - - -def _static(arg, name: str, what: str, axis: int) -> int: - """One extent as the compile-time count the program dimensions need.""" - value = static_dim_value(arg) - if value is None: - raise ValueError( - f"launch of {name!r}: {what}[{axis}] is decided at run time, and " - f"only a launch-provided CTA count can be" - ) - return value - - -__all__ = [ - "Geometry", - "coarsest_topology", - "launch_geometry", - "topology_domains", -] diff --git a/src/tilefoundry/compile.py b/src/tilefoundry/compile.py index 306db7bc5..3904a9100 100644 --- a/src/tilefoundry/compile.py +++ b/src/tilefoundry/compile.py @@ -10,12 +10,10 @@ import tempfile from dataclasses import dataclass, replace +from tilefoundry.codegen.context_builder import build_codegen_context from tilefoundry.codegen.cpu.context import CpuCodegenContext from tilefoundry.codegen.cuda.context import CudaCodegenContext from tilefoundry.codegen.linker import link_modules -from tilefoundry.codegen.registry import group_functions_by_target -from tilefoundry.codegen.signature import symbol_table -from tilefoundry.codegen.topology import launch_geometry, topology_domains from tilefoundry.inspection import as_script as _as_script from tilefoundry.ir.core.module import Module from tilefoundry.ir.hir.function import Function as HirFunction @@ -102,10 +100,7 @@ def build( f"conflicts with the Module Target {_target_summary(module_target)}" ) - mod = _build_default_pipeline().run(mod) - workdir = os.path.join( - tempfile.gettempdir(), f"tilefoundry_build_{mod.entry}_{os.getpid()}_split" - ) + workdir = tempfile.mkdtemp(prefix=f"tilefoundry_build_{mod.entry}_split_") return _build_split_runtime_module(mod, workdir=workdir) @@ -120,69 +115,64 @@ def _build_split_runtime_module(mod: Module, *, workdir: str) -> "RuntimeModule" load each other's. Unsupported module shapes raise during codegen -- there is no fallback to a single-source path. """ - device_target = _device_target(mod) + codegen = build_codegen_context(mod) cpu_entry = mod.entry_function() if not isinstance(cpu_entry.target, CpuTarget): raise ValueError( - f"tilefoundry.build: entry {cpu_entry.name!r} is not a CPU host entry " - f"after normalization" + f"tilefoundry.build: entry {cpu_entry.name!r} is not a CPU host entry; " + "build() accepts pass-prepared TIR" ) - symbols = symbol_table(mod, device_target) - launches = launch_geometry(mod) - device_modules = tuple( - device_target.get_code_generator().emit( - domain, - device_fns, - device_target, - CudaCodegenContext(symbols=symbols, target=device_target, launches=launches), - ) - for domain, device_fns in _device_domains(topology_domains(mod), device_target) - ) - host_module = cpu_entry.target.get_code_generator().emit( - mod, - (cpu_entry,), - cpu_entry.target, - CpuCodegenContext(symbols=symbols, target=cpu_entry.target), - ) + device_modules = [] + host_module = None + for group in codegen.groups: + view = codegen.for_group(group) + if isinstance(group.target, CudaTarget): + context = CudaCodegenContext( + symbols=view.symbols, + target=view.target, + launches=view.launches, + codegen_context=view.codegen, + ) + device_modules.append( + group.target.get_code_generator().emit( + group.owner, group.functions, group.target, context + ) + ) + elif isinstance(group.target, CpuTarget): + if cpu_entry not in group.functions: + continue + context = CpuCodegenContext( + symbols=view.symbols, + target=view.target, + resolved_launches=view.resolved_launches, + codegen_context=view.codegen, + ) + host_module = group.target.get_code_generator().emit( + group.owner, (cpu_entry,), group.target, context + ) + if host_module is None: + raise ValueError("tilefoundry.build: root context produced no host entry group") units = (*device_modules, host_module) digest = hashlib.sha256("".join(m.source for m in units).encode("utf-8")).hexdigest()[:16] - loaded_as = replace(symbols[id(cpu_entry)], name=cpu_entry.name) + entry_signature = codegen.symbols.by_function[id(cpu_entry)].callable + loaded_as = replace(entry_signature, name=cpu_entry.name) + device_targets = tuple( + group.target for group in codegen.groups if isinstance(group.target, CudaTarget) + ) + if not device_targets: + raise ValueError(f"tilefoundry.build: module {mod.name!r} has no CUDA device functions") linked_module = link_modules( units, workdir=os.path.join(workdir, digest), lib_name=cpu_entry.name, entry=loaded_as, - cuda_arch=device_target.arch.removeprefix("sm_"), + cuda_arch=device_targets[0].arch.removeprefix("sm_"), ) return load_linked_module(linked_module) -def _device_target(mod: Module) -> Target: - """The one device Target this module's functions run on. - - A second unequal one is refused by the grouping that sees them both, - because one linked artifact holds one device architecture. - """ - targets = [t for t in group_functions_by_target(mod) if isinstance(t, CudaTarget)] - if not targets: - raise ValueError(f"tilefoundry.build: module {mod.name!r} has no CUDA device functions") - return targets[0] - - -def _device_domains(domains, device_target: Target): - """Each *domains* entry that has device functions, and which of them they are. - - A module states its instance counts once, so one that states its own is one - translation unit; a domain holding only host functions is not one at all. - """ - for domain, functions in domains: - device_fns = tuple(fn for fn in functions if fn.target == device_target) - if device_fns: - yield domain, device_fns - - def compile( mod: Module, /, @@ -193,7 +183,8 @@ def compile( *mod* must be a ``Module``. Meshes are derived from the IR body. """ - return build(mod, target=target) + prepared = _build_default_pipeline().run(mod) + return build(prepared, target=target) def _canonical_module_text(mod: Module) -> str: diff --git a/src/tilefoundry/inspection/values.py b/src/tilefoundry/inspection/values.py index d78453163..79d8cc8bb 100644 --- a/src/tilefoundry/inspection/values.py +++ b/src/tilefoundry/inspection/values.py @@ -1,304 +1,34 @@ -"""How a record renders as one line of comment. - -A record states typed fields ([core-ir §2](docs/spec/core-ir.md#2-expr)); what -they look like is decided here, so a family arrives by declaring what it has -rather than by writing a form string -- which is how five came to spell one value -five ways. A value earns a rendering of its own only where Python cannot already -read it: an ``int``, a token, and a mapping of them get none. The constants below -are the ladder of [inspection §2.8](docs/spec/inspection.md#28-record-comment-forms), -one each, and nothing else in the tree spells one. -""" +"""Canonical metadata comment printer.""" from __future__ import annotations import json -from collections.abc import Callable, Mapping -from dataclasses import dataclass, field, fields -from typing import get_type_hints - -from tilefoundry.analysis.metadata import ( - Breakdown, - ComputeCostMetadata, - LoopFootprintMetadata, - MemoryMetadata, - PerformanceMetadata, - PerformanceSummaryMetadata, - RooflineMetadata, - TrafficMetadata, -) -from tilefoundry.analysis.report import ( - declare_record, - declared_records, - expr_field, - expr_fields, - family_of, - render_record, +from collections.abc import Mapping +from dataclasses import dataclass, field, fields, is_dataclass + +from tilefoundry.ir.core.metadata import IRMetadata +from tilefoundry.ir.core.values import TripInterval + +PAIR, PER_UNIT, ENTRY, ENTRIES, FIELD, FIELDS, PARTS, TRIPS = ( + "/", + "@", + ":", + ",", + "=", + " ", + "; ", + "*", ) -from tilefoundry.ir.core.metadata import IRMetadata, SourceSpanMetadata -from tilefoundry.ir.core.values import TotalAndPerUnit, TripInterval -from tilefoundry.visitor_registry.contexts import TrafficBytes - -PAIR = "/" -PER_UNIT = "@" -ENTRY = ":" -ENTRIES = "," -FIELD = "=" -FIELDS = " " -PARTS = "; " -TRIPS = "*" - - -class Prose(str): - """Text that is a sentence rather than a token. - - A token renders bare, which only works while it holds no separator. A - sentence holds several -- spaces, commas, a semicolon -- so it renders as a - quoted, escaped string literal, which brackets it: a reader splits a layer - outside the quotes ([inspection §2.8](docs/spec/inspection.md#28-record-comment-forms)). - """ - - -def _trip_interval(value: TripInterval, render: Callable[[object], str]) -> str: - """One interval, offset by the trip index when it repeats.""" - if value.trips <= 1: - return f"[{value.start}{ENTRIES}{value.end})" - offset = f"{value.stride}t+" - return f"[{offset}{value.start}{ENTRIES}{offset}{value.end}){TRIPS}{value.trips}" - - -RENDER: dict[type, Callable[..., str]] = { - Prose: lambda value, render: json.dumps(str(value)), - TrafficBytes: lambda value, render: f"r{value.read}{PAIR}w{value.write}", - TotalAndPerUnit: ( - lambda value, render: f"{render(value.total)}{PER_UNIT}{render(value.per_unit)}" - ), - TripInterval: _trip_interval, -} - - -def render_value(value: object) -> str: - """One value, rendered by its own type or as the entries it holds.""" - for value_type, render in RENDER.items(): - if isinstance(value, value_type): - return render(value, render_value) - entries = _entries_of(value) - if entries is not None: - return ENTRIES.join(f"{key}{ENTRY}{render_value(item)}" for key, item in entries) - if isinstance(value, tuple | list): - return ENTRIES.join(render_value(item) for item in value) - return str(value) - - -def _entries_of(value: object) -> tuple[tuple[object, object], ...] | None: - """*value* as key/value pairs, or ``None`` when it is not a mapping.""" - if isinstance(value, Mapping): - return tuple(value.items()) - if isinstance(value, tuple) and all( - isinstance(item, tuple) and len(item) == 2 for item in value - ): - return value - return None - _UNSET = object() -@dataclass(frozen=True) -class Projection: - """One key a comment emits: its name, its type, and where its value is read. - - A field is the identity projection of itself. Anything else -- a sum, a - count, several fields folded into one value -- is declared here, which is - what keeps the next one from appearing unannounced. - - *default* is the value this key says nothing by, so it is left out. Without - one, an empty mapping says nothing and every other value is emitted. - """ - - key: str - type: object - of: Callable[..., object] - default: object = _UNSET - opt_in: bool = False - - -@dataclass(frozen=True) -class RecordComment: - """What one record type emits, in order, under which family name.""" - - family: str - emissions: tuple[Projection, ...] - - -_COMMENTS: dict[type[IRMetadata], RecordComment] = {} - - -def _field_projections(record: type[IRMetadata], names: tuple[str, ...]) -> tuple[Projection, ...]: - hints = get_type_hints(record) - defaults = {item.name: item.default for item in fields(record)} - return tuple(Projection(name, hints[name], _read(name), defaults[name]) for name in names) - - -def _read(name: str) -> Callable[..., object]: - return lambda record: getattr(record, name) - - -def comment( - record: type[IRMetadata], - *emitted: str | Projection, - family: str | None = None, -) -> None: - """Declare that *record* renders as a comment, and what it emits. - - A field is named by its own name; anything else is a ``Projection``. Naming - nothing emits every field in declaration order. A record with no declaration - renders as ``None``, which is how metadata that is not a report stays out of - the comment. - - *family* is only for a record whose reported name is not the one its class - name states. - """ - declared = emitted or tuple(item.name for item in fields(record)) - emissions = tuple( - item if isinstance(item, Projection) else _field_projections(record, (item,))[0] - for item in declared - ) - declare_record(record, family=family) - _COMMENTS[record] = RecordComment(family=family_of(record), emissions=emissions) - - -def comment_of(record: type[IRMetadata]) -> RecordComment | None: - """What *record* declared it emits, if it declared anything.""" - return _COMMENTS.get(record) - - -def render_comment(record: IRMetadata, *, opt_in: frozenset[str] = frozenset()) -> str | None: - """One record as one comment, or ``None`` when it does not report. - - A key is its declared name with ``_`` written as ``-``; a unit belongs in - that name, which is where it is said once. A record declaring one key uses - the family name as that key, because for a record of one thing the family - and the key say the same thing twice. - """ - declared = _COMMENTS.get(type(record)) - if declared is None: - return None - alone = len(declared.emissions) == 1 - emitted: list[str] = [] - for emission in declared.emissions: - if emission.opt_in and emission.key not in opt_in: - continue - value = emission.of(record) - if _says_nothing(value, emission.default): - continue - key = declared.family if alone else emission.key.replace("_", "-") - emitted.append(f"{key}{FIELD}{render_value(value)}") - if alone and emitted: - return emitted[0] - return FIELDS.join([declared.family, *emitted]) - - -def _says_nothing(value: object, default: object) -> bool: - """Whether this value adds nothing to the comment it would appear in.""" - if default is not _UNSET: - return value == default - entries = _entries_of(value) - return entries is not None and not entries - - -def _paired[V]( - held: "Breakdown[V]", topologies: tuple[str, ...] -) -> dict[str, TotalAndPerUnit[dict[str, V]]]: - """Each kind's quantity, whole beside what one unit of each level holds. - - The level names come from the record's own ``topologies``, which is where - they are written once; a ``Spread`` states its shares in that order and - carries no names of its own. - """ - return { - kind: TotalAndPerUnit(spread.total, dict(zip(topologies, spread.per_unit, strict=False))) - for kind, spread in held.kinds - } - - -def _paired_flops(record: ComputeCostMetadata): - """Each dtype's work, whole and per unit at every declared level.""" - return _paired(record.flops, record.topologies) - - -def _paired_service(record: ComputeCostMetadata): - """Each service kind's work, whole and per unit at every declared level. - - What a machine is asked for that is not floating point: comparing, selecting, - integer arithmetic, a reciprocal, a local move. Reported beside the flops - rather than folded into them, because a predicate priced as a FLOP is a - number about a pipe the work never went down. - """ - return _paired(record.service, record.topologies) - - -def _paired_traffic(record: TrafficMetadata): - """Each storage level's traffic, whole and per unit at every declared level.""" - return _paired(record.storage, record.topologies) - - -def _by_operand(record: TrafficMetadata) -> dict[str, TrafficBytes]: - """What each operand moved, positional against ``(*call.args, call)``.""" - last = len(record.operands) - 1 - return { - ("result" if index == last else str(index)): moved - for index, moved in enumerate(record.operands) - } - - -def peak_footprint(record: MemoryMetadata) -> dict[str, int]: - """How much of each level the function holds at its peak.""" - return {item.memory_level: item.peak_bytes for item in record.footprint} - - -def _persistent_bytes(record: MemoryMetadata) -> int: - """The part of the peak the function cannot reclaim, across levels.""" - return sum(item.persistent_bytes for item in record.footprint) - - -def _advisory_count(record: MemoryMetadata) -> int: - """How many advisories there are; the comment is not where they are read.""" - return len(record.advisories) - - -def _loop_footprints(record: LoopFootprintMetadata) -> dict[str, str]: - return { - f"{item.buffer}@{item.memory_level}": ( - f"{item.bytes}/{item.device_bytes}/{item.repeated_bytes}" - ) - for item in record.footprints - } - - -def _loop_footprint_status(record: LoopFootprintMetadata) -> str: - return "complete" if record.known else "lower-bound" - - -def _interval(record: PerformanceMetadata) -> TripInterval: - """The occurrence's interval, with its repetition folded in.""" - timeline = record.timeline - return TripInterval(timeline.start_ns, timeline.end_ns, timeline.stride_ns, timeline.trips) - - -def _predicted_ns(record: PerformanceSummaryMetadata) -> int: - """How long the whole Function is predicted to take.""" - return record.timeline.end_ns - record.timeline.start_ns - - -def _source_span(record: SourceSpanMetadata) -> str: - """Where the expression was authored, as one location.""" - return f"{record.file}:{record.line}:{record.column}" +class Prose(str): + """Sentence text rendered as a quoted DSL string.""" @dataclass(frozen=True) class ReportIdentity(IRMetadata): - """Which program, on which machine, this report is about.""" - target: str = "" module: str = "" function: str = "" @@ -307,120 +37,225 @@ class ReportIdentity(IRMetadata): @dataclass(frozen=True) class ReportSelection(IRMetadata): - """What was asked for, and what running it took.""" - requested: tuple[str, ...] = () executed: tuple[str, ...] = () @dataclass(frozen=True) class MemorySummary(IRMetadata): - """What one function holds at its peak, per level.""" - peak_bytes: dict[str, int] = field(default_factory=dict) @dataclass(frozen=True) class AdvisorySummary(IRMetadata): - """One thing the memory walk observed that a reader should weigh.""" - text: Prose @dataclass(frozen=True) class PerformanceSummaryView(IRMetadata): - """One function's prediction, under the root the report is about. - - ``root`` is the report's own identity composed for a reader, not something - the family measured, which is why it lives here and not on - ``PerformanceSummaryMetadata``. ``waves`` is stated even when it is one: how - many passes over the machine a plan takes is a conclusion, and one wave is an - answer rather than nothing to say. It is declared with no value it says - nothing by, so the suppression rule itself stays one rule. - """ - root: str = "" predicted_ns: int = 0 waves: int = 1 -comment( - ComputeCostMetadata, - Projection("flops", dict[str, TotalAndPerUnit[dict[str, int]]], _paired_flops), - Projection("service", dict[str, TotalAndPerUnit[dict[str, int]]], _paired_service), -) -comment( - TrafficMetadata, - Projection( - "traffic", - dict[str, TotalAndPerUnit[dict[str, TrafficBytes]]], - _paired_traffic, - ), - Projection("operands", dict[str, TrafficBytes], _by_operand, opt_in=True), -) -comment( - MemoryMetadata, - Projection("peak", dict[str, int], peak_footprint), - Projection("persistent", int, _persistent_bytes, default=0), - Projection("advisories", int, _advisory_count, default=0), -) -comment( - LoopFootprintMetadata, - Projection("footprints", dict[str, str], _loop_footprints), - Projection("status", str, _loop_footprint_status), -) -comment(RooflineMetadata, "ideal_ns", "bound_by") -comment( - PerformanceMetadata, - Projection("interval", TripInterval, _interval), - family="performance", -) -comment( - PerformanceSummaryMetadata, - Projection("predicted_ns", int, _predicted_ns), - Projection("waves", int, _read("waves")), - family="performance", -) -comment(SourceSpanMetadata, Projection("span", str, _source_span), family="source") -comment(ReportIdentity, family="analysis") -comment(ReportSelection, family="selection") -comment(MemorySummary, family="peak-footprint") -comment(AdvisorySummary, family="advisory") -comment( - PerformanceSummaryView, - "root", - "predicted_ns", - Projection("waves", int, _read("waves")), - family="performance", -) +class CommentPrinter: + """Render metadata by explicit ``print_`` methods.""" + + def print(self, value): + method = getattr(self, f"print_{type(value).__name__}", None) + if method: + return method(value) + if is_dataclass(value): + return ENTRIES.join( + f"{item.name}{ENTRY}{self.print(getattr(value, item.name))}" + for item in fields(value) + ) + if isinstance(value, Mapping): + return ";".join(f"{k}{ENTRY}{self.print(v)}" for k, v in value.items()) + if isinstance(value, (tuple, list)): + return ENTRIES.join(self.print(v) for v in value) + return str(value) + + def print_Prose(self, value): + return json.dumps(str(value)) + + def print_TrafficBytes(self, value): + return f"r{value.read}{PAIR}w{value.write}" + + def print_TripInterval(self, value): + if value.trips <= 1: + return f"[{value.start},{value.end})" + offset = f"{value.stride}t+" + return f"[{offset}{value.start},{offset}{value.end}){TRIPS}{value.trips}" + + @staticmethod + def _empty_without_default(value): + if isinstance(value, Mapping | tuple | list): + return not value + return False + + def _record(self, family, values): + out = [] + for item in values: + key, value = item[:2] + default = item[2] if len(item) == 3 else _UNSET + if default is not _UNSET: + if value == default: + continue + elif self._empty_without_default(value): + continue + out.append(f"{key.replace('_', '-')}={self.print(value)}") + return FIELDS.join([family, *out]) if out else family + + def _single(self, family, value, default=_UNSET): + if default is not _UNSET and value == default: + return family + if default is _UNSET and self._empty_without_default(value): + return family + return f"{family}{FIELD}{self.print(value)}" + + def _spread(self, spread, topologies, *, logical=True): + values = [] + if logical: + values.append(f"{self.print(spread.logical)}@logical") + values.append(f"{self.print(spread.total)}@total") + values.extend( + f"{self.print(value)}@{topology}" + for topology, value in zip(topologies, spread.per_unit, strict=False) + ) + return self.print(values) + + def _breakdown(self, held, topologies, *, logical=True): + return { + kind: self._spread(spread, topologies, logical=logical) + for kind, spread in held.kinds + } + + def print_ComputeCostMetadata(self, record, **_): + return self._record( + "compute-cost", + ( + ("flops", self._breakdown(record.flops, record.topologies)), + ("other_ops", self._breakdown(record.other_ops, record.topologies)), + ), + ) + + def print_TrafficMetadata(self, record, *, opt_in=frozenset()): + traffic = self._breakdown(record.storage, record.topologies, logical=False) + values = [("traffic", traffic)] + if "operands" in opt_in: + last = len(record.operands) - 1 + operands = { + "result" if index == last else str(index): moved + for index, moved in enumerate(record.operands) + } + values.append(("operands", operands)) + return self._record("traffic", values) + + def print_MemoryMetadata(self, record, **_): + return self._record( + "memory", + ( + ("peak", {item.memory_level: item.peak_bytes for item in record.footprint}), + ("persistent", sum(item.persistent_bytes for item in record.footprint), 0), + ("advisories", len(record.advisories), 0), + ), + ) + + def print_LoopFootprintMetadata(self, record, **_): + footprints = { + f"{item.buffer}@{item.memory_level}": PAIR.join( + str(value) for value in (item.bytes, item.device_bytes, item.repeated_bytes) + ) + for item in record.footprints + } + return self._record( + "loop-footprint", + (("footprints", footprints), ("status", "complete" if record.known else "lower-bound")), + ) + + def print_RooflineMetadata(self, record, **_): + return self._record( + "roofline", (("ideal_ns", record.ideal_ns, 0), ("bound_by", record.bound_by, "none")) + ) + + def print_PerformanceMetadata(self, record, **_): + interval = TripInterval( + record.timeline.start_ns, + record.timeline.end_ns, + record.timeline.stride_ns, + record.timeline.trips, + ) + return self._single("performance", interval) + + def print_PerformanceSummaryMetadata(self, record, **_): + predicted = record.timeline.end_ns - record.timeline.start_ns + return self._record("performance", (("predicted_ns", predicted), ("waves", record.waves))) + + def print_SourceSpanMetadata(self, record, **_): + return self._single("source", f"{record.file}:{record.line}:{record.column}") + + def print_ReportIdentity(self, record, **_): + return self._record( + "analysis", + ( + ("target", record.target, ""), + ("module", record.module, ""), + ("function", record.function, ""), + ("topology", record.topology, "none"), + ), + ) + + def print_ReportSelection(self, record, **_): + return self._record( + "selection", (("requested", record.requested, ()), ("executed", record.executed, ())) + ) + + def print_MemorySummary(self, record, **_): + return self._single("peak-footprint", record.peak_bytes, {}) + + def print_AdvisorySummary(self, record, **_): + return self._single("advisory", record.text) + + def print_PerformanceSummaryView(self, record, **_): + return self._record( + "performance", + ( + ("root", record.root, ""), + ("predicted_ns", record.predicted_ns, 0), + ("waves", record.waves), + ), + ) + + +_PRINTER = CommentPrinter() + + +def render_comment(record, *, opt_in=frozenset()): + method = getattr(_PRINTER, f"print_{type(record).__name__}", None) + return method(record, opt_in=opt_in) if method else None + + +def peak_footprint(record): + return {item.memory_level: item.peak_bytes for item in record.footprint} __all__ = [ - "AdvisorySummary", - "ENTRIES", + "CommentPrinter", + "Prose", + "render_comment", + "PAIR", + "PER_UNIT", "ENTRY", + "ENTRIES", "FIELD", "FIELDS", - "PAIR", "PARTS", - "PER_UNIT", - "Prose", - "RENDER", "TRIPS", - "MemorySummary", - "PerformanceSummaryView", - "Projection", - "RecordComment", "ReportIdentity", "ReportSelection", - "comment", - "comment_of", - "declared_records", - "expr_field", - "expr_fields", - "family_of", + "MemorySummary", + "AdvisorySummary", + "PerformanceSummaryView", "peak_footprint", - "render_comment", - "render_record", - "render_value", ] diff --git a/src/tilefoundry/passes/transforms/host_entry.py b/src/tilefoundry/passes/transforms/host_entry.py index b385a944b..f0d353c55 100644 --- a/src/tilefoundry/passes/transforms/host_entry.py +++ b/src/tilefoundry/passes/transforms/host_entry.py @@ -11,16 +11,11 @@ from dataclasses import replace -from tilefoundry.ir.core import Call, Constant, Tuple, Var +from tilefoundry.ir.core import Var from tilefoundry.ir.core.module import Module -from tilefoundry.ir.hir.loop_region import LoopRegion from tilefoundry.ir.tir.launch import launch_call from tilefoundry.ir.tir.prim_function import PrimFunction -from tilefoundry.ir.tir.shape import ShapeOf -from tilefoundry.ir.tir.stmts import LetStmt, MeshScope, Sequential -from tilefoundry.ir.tir.symbol_ref import SymbolRef -from tilefoundry.ir.types.shard.shard_layout import ShardLayout -from tilefoundry.ir.visitor import ExprWalker +from tilefoundry.ir.tir.stmts import Sequential from tilefoundry.passes.pass_base import ModulePass from tilefoundry.target import CpuTarget, CudaTarget @@ -64,130 +59,26 @@ def _device_callee(module: Module) -> PrimFunction: return device_fns[0] -def _launch_config_from_body( - body: Sequential, -) -> tuple[tuple[int | None, int, int], tuple[int, int, int]]: - """Derive grid and block dimensions from body mesh topologies. - - CTA topology sizes multiply into ``grid.x`` and thread sizes into - ``block.x``; warp axes do not contribute. A launch-provided CTA extent - yields ``grid.x = None`` for static callers to reject. Other dimensions - remain one until a public convention exists. - """ - grid_x = 1 - block_x = 1 - cta_dynamic = False - - def _topo_dims(mesh) -> tuple[int, int]: - """Return grid and block contributions for *mesh*. - - Return ``(grid_size_contribution, block_size_contribution)`` - for *mesh*'s full topology list. - """ - nonlocal cta_dynamic - g = 1 - b = 1 - for t in mesh.topologies: - tname = t.name - size = t.size - if not isinstance(size, int): - if tname == "cta": - cta_dynamic = True - continue - raise ValueError( - f"_launch_config_from_body: topology {tname!r} has a " - f"dynamic/scalar extent ({size!r}) that cannot be converted " - f"to a static launch config; only a 'cta' level may be " - f"launch-provided" - ) - tsize = size - if tname == "cta": - g *= tsize - else: - b *= tsize - return g, b - - def _harvest_from_layout(layout) -> None: - nonlocal grid_x, block_x - if isinstance(layout, ShardLayout): - g, b = _topo_dims(layout.mesh) - grid_x = max(grid_x, g) - block_x = max(block_x, b) - - class _LayoutVisitor(ExprWalker[None]): - def _record(self, expr) -> None: - _harvest_from_layout(getattr(getattr(expr, "type", None), "layout", None)) - - def visit_Call(self, expr: Call, ctx=None) -> None: - self._record(expr) - self.visit_operands(expr, ctx) - - def visit_Tuple(self, expr: Tuple, ctx=None) -> None: - self._record(expr) - - def visit_LoopRegion(self, expr: LoopRegion, ctx=None) -> None: - self._record(expr) - - def visit_Var(self, expr: Var, ctx=None) -> None: - self._record(expr) - - def visit_Constant(self, expr: Constant, ctx=None) -> None: - self._record(expr) - - def visit_SymbolRef(self, expr: SymbolRef, ctx=None) -> None: - self._record(expr) - - def visit_ShapeOf(self, expr: ShapeOf, ctx=None) -> None: - self._record(expr) - - def default_visit(self, expr, ctx=None) -> None: - self._record(expr) - - def walk(stmt) -> None: - nonlocal grid_x, block_x - match stmt: - case MeshScope(): - g, b = _topo_dims(stmt.mesh) - grid_x = max(grid_x, g) - block_x = max(block_x, b) - walk(stmt.body) - case Sequential(): - for s in stmt.body: - walk(s) - case LetStmt(): - if hasattr(stmt, "value"): - _LayoutVisitor().visit(stmt.value) - if hasattr(stmt, "var") and getattr(stmt.var, "type", None) is not None: - _harvest_from_layout(getattr(stmt.var.type, "layout", None)) - walk(stmt.body) - - walk(body) - grid = (None, 1, 1) if cta_dynamic else (grid_x, 1, 1) - return grid, (block_x, 1, 1) +def _product(values: tuple[object, ...]) -> object: + result: object = 1 + for value in values: + if value is None: + raise ValueError( + "insert_default_host_entry: topology extent cannot be None; " + "declare a DimVar-backed host-computable extent" + ) + result = result * value + return result def _launch_config_of( - fn: PrimFunction, -) -> tuple[tuple[int | None, int, int], tuple[int, int, int]]: - """The geometry one launch of *fn* runs at. - - A specialization prototype has variants and no body of its own, so the - variants state it -- and they have to agree, because a translation unit - states its program dimensions once for every kernel it holds. - """ - if fn.variants and not fn.body.body: - by_variant = { - variant.name: _launch_config_from_body(variant.body) for variant in fn.variants - } - distinct = set(by_variant.values()) - if len(distinct) > 1: - raise ValueError( - f"insert_default_host_entry: prototype {fn.name!r} has variants " - f"launching at different geometries ({by_variant}); one " - f"translation unit runs every kernel in it at one geometry" - ) - return distinct.pop() - return _launch_config_from_body(fn.body) + module: Module, +) -> tuple[tuple[object, object, object], tuple[object, object, object]]: + """Derive the implicit Launch expressions solely from module topologies.""" + topologies = module.effective_topologies() + grid = _product(tuple(t.size for t in topologies if t.name == "cta")) + block = _product(tuple(t.size for t in topologies if t.name == "thread")) + return (grid, 1, 1), (block, 1, 1) def insert_default_host_entry(module: Module) -> Module: @@ -203,13 +94,7 @@ def insert_default_host_entry(module: Module) -> Module: device_fn = _device_callee(module) entry_params = tuple(Var(type=p.type, name=p.name) for p in device_fn.params) - grid, block = _launch_config_of(device_fn) - if grid[0] is None: - raise ValueError( - f"insert_default_host_entry: device function {device_fn.name!r} has " - f"a launch-provided (dynamic) CTA extent; the implicit host entry " - f"cannot derive its grid — launch it from an explicit host entry" - ) + grid, block = _launch_config_of(module) launch = launch_call(device_fn, entry_params, grid, block) name = _fresh_entry_name({fn.name for fn in module.functions}) entry = PrimFunction( diff --git a/src/tilefoundry/visitor_registry/contexts.py b/src/tilefoundry/visitor_registry/contexts.py index 3705490bd..f51c5fc8f 100644 --- a/src/tilefoundry/visitor_registry/contexts.py +++ b/src/tilefoundry/visitor_registry/contexts.py @@ -193,6 +193,11 @@ class Cost: name instead, from the sending unit's own view. """ + @property + def ops(self) -> Mapping[str, int]: + """Program-work spelling for non-floating-point operation counts.""" + return self.service + @property def bytes(self) -> int: """Every operand's traffic, in either direction.""" diff --git a/tests/analysis/test_analysis_families.py b/tests/analysis/test_analysis_families.py index 0da3e23a7..2109a3e24 100644 --- a/tests/analysis/test_analysis_families.py +++ b/tests/analysis/test_analysis_families.py @@ -14,23 +14,22 @@ import pytest -from tests.fixtures.placed.fused_boundary import FusedBoundary from tests.fixtures.placed.performance_findings import _CrossScopePerformance from tests.fixtures.placed.symbolic_offset import ( _LiteralStoreOffset, _SymbolicStoreOffset, ) -from tests.fixtures.placed.tp_all_to_all import GPUS, SENT_BYTES, TransposeShard from tilefoundry import func, module from tilefoundry.analysis import ( + Breakdown, ComputeCostMetadata, MemoryHierarchyFacts, MemoryMetadata, - ParallelCapacityFacts, PerformanceMetadata, PerformanceServiceFacts, PerformanceSummaryMetadata, RooflineMetadata, + Spread, ThroughputFacts, TrafficMetadata, ) @@ -39,7 +38,6 @@ _local_duration_ns, ) from tilefoundry.analysis.errors import AnalysisError -from tilefoundry.analysis.metadata import Breakdown, Spread, shares from tilefoundry.dsl import ConstTensor, DimVar, Mesh, Tensor, Topology, tf from tilefoundry.ir.core import ( Call, @@ -235,9 +233,11 @@ def test_a_symbolic_store_stride_preserves_the_literal_control_result() -> None: bound = get_metadata(result.function, RooflineMetadata) summary = get_metadata(result.function, PerformanceSummaryMetadata) assert cost is not None and bound is not None and summary is not None + integer_ops = cost.other_ops.of("integer") + assert integer_ops is not None observed[name] = ( - shares(cost.service, cost.topologies)["integer"], - shares(cost.service, cost.topologies, "cta")["integer"], + integer_ops.total, + integer_ops.per_unit[0], bound.ideal_ns, summary.timeline.end_ns - summary.timeline.start_ns, ) @@ -323,9 +323,11 @@ def test_a_matmul_counts_its_rows_once_whichever_axis_the_mesh_split() -> None: ) cost = get_metadata(product, ComputeCostMetadata) summary = get_metadata(report.function, PerformanceSummaryMetadata) + bf16 = cost.flops.of("bf16") + assert bf16 is not None per_layout[name] = ( - shares(cost.flops, cost.topologies)["bf16"], - shares(cost.flops, cost.topologies, "cta")["bf16"], + bf16.total, + bf16.per_unit[0], summary.timeline.end_ns, ) @@ -359,7 +361,7 @@ def test_a_program_whose_buffers_have_nowhere_to_sit_is_refused() -> None: held = get_metadata( analyze(_SharedTile, unrestated, analysis="memory").function, MemoryMetadata ).footprint - assert next(item.peak_bytes for item in held if item.memory_level == "smem") == 422_400 + assert next(item.peak_bytes for item in held if item.level == "smem") == 422_400 fits = analyze( roomy, @@ -378,10 +380,10 @@ def test_a_program_whose_buffers_have_nowhere_to_sit_is_refused() -> None: analysis="memory", ) assert [ - (item.binding, item.memory_level, item.bytes, item.defined_at, item.last_used_at) + (item.binding, item.level, item.bytes, item.defined_at, item.last_used_at) for item in get_metadata(fits.function, MemoryMetadata).lifetimes ] == [ - (item.binding, item.memory_level, item.bytes, item.defined_at, item.last_used_at) + (item.binding, item.level, item.bytes, item.defined_at, item.last_used_at) for item in get_metadata(relieved.function, MemoryMetadata).lifetimes ] @@ -405,7 +407,8 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N work = next( record for record in records - if record is not None and any(shares(record.flops, record.topologies, "cta").values()) + if record is not None + and any(spread.per_unit[0] for _name, spread in record.flops.kinds) ) with pytest.raises( @@ -413,18 +416,17 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N match=r"^performance: selected topology level 'thread', but the target's " r"one-unit throughputs are stated for 'cta'$", ): - _local_duration_ns(work, throughput, services, topology_level="thread") + _local_duration_ns(work, throughput, services, level="thread") with pytest.raises(AnalysisError, match=r"unknown compute dtype 'f9e9m9'"): _local_duration_ns( replace( work, - topologies=("cta",), - flops=Breakdown((("f9e9m9", Spread(8, (8,))),)), + flops=Breakdown((*work.flops.kinds, ("f9e9m9", Spread(8, 8, (8,))))), ), throughput, services, - topology_level="cta", + level="cta", ) crossed = TrafficMetadata( @@ -433,7 +435,11 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N ( ( throughput.bandwidth_level, - Spread(TrafficBytes(read=4096), (TrafficBytes(read=4096),)), + Spread( + TrafficBytes(read=4096), + TrafficBytes(read=4096), + (TrafficBytes(read=4096),), + ), ), ) ), @@ -447,73 +453,5 @@ def test_a_price_is_refused_where_the_machine_states_no_rate_to_pay_it_at() -> N throughput, replace(services, unit_bandwidth=()), moved=crossed, - topology_level="cta", + level="cta", ) - - -def test_one_pass_states_the_work_at_every_level_and_prices_it_there() -> None: - """The numbers this program reports at each level it declares, not the relations. - - One analysis states one unit's work at every declared level, and what one - unit gets through is the device peak over however many of that unit the - device holds, so a thread is priced far slower than a CTA for the same - share. Nothing here is sharded across ``thread``, so a thread repeats its - CTA's share rather than dividing it: asking at ``thread`` unrolls that - repetition into the whole-program total and leaves the shares alone. - """ - function = FusedBoundary.entry_function() - reported = {} - for level in ("cta", "thread"): - result = analyze(FusedBoundary, function, analysis=("performance",), topology_level=level) - cost = get_metadata(result.function, ComputeCostMetadata) - summary = get_metadata(result.function, PerformanceSummaryMetadata) - reported[level] = { - "answered_at": result.topology_level, - "levels": cost.topologies, - "whole_program": shares(cost.flops, cost.topologies), - "per_cta": shares(cost.flops, cost.topologies, "cta"), - "per_thread": shares(cost.flops, cost.topologies, "thread"), - "end_ns": summary.timeline.end_ns, - } - - assert reported == { - "cta": { - "answered_at": "cta", - "levels": ("cta", "thread"), - "whole_program": {"bf16": 576, "f32": 128}, - "per_cta": {"bf16": 72, "f32": 16}, - "per_thread": {"bf16": 72, "f32": 16}, - "end_ns": 6, - }, - "thread": { - "answered_at": "thread", - "levels": ("cta", "thread"), - "whole_program": {"bf16": 576, "f32": 2112}, - "per_cta": {"bf16": 72, "f32": 16}, - "per_thread": {"bf16": 72, "f32": 16}, - "end_ns": 1849, - }, - } - capacity = FusedBoundary.resolve_target().get_facts(ParallelCapacityFacts, "gpu") - assert capacity.parallel_units == 1, "a target that never said how many cards is one card" - - -def test_a_reshard_that_changes_shards_says_what_left_the_card() -> None: - """Splitting one tensor two ways is an exchange, and it costs both resources. - - What crosses is counted against the boundary it crossed, which no storage - level names: the data is global memory at both ends and has still gone to - another card. A unit inside the boundary is not one of the parties, so it - states no share of the crossing. - """ - result = analyze( - TransposeShard, TransposeShard.entry_function(), analysis=("compute-cost", "memory") - ) - moved = get_metadata(result.function, TrafficMetadata) - - crossing = moved.communication.of("gpu") - assert crossing.total == TrafficBytes(read=SENT_BYTES * GPUS, write=SENT_BYTES * GPUS) - by_level = dict(zip(moved.topologies, crossing.per_unit, strict=True)) - assert by_level["gpu"] == TrafficBytes(read=SENT_BYTES, write=SENT_BYTES) - assert by_level["cta"] == TrafficBytes() - assert moved.communication.of("cta") is None diff --git a/tests/analysis/test_analyze_at_a_size.py b/tests/analysis/test_analyze_at_a_size.py index 56c13dcb4..e1a3a106e 100644 --- a/tests/analysis/test_analyze_at_a_size.py +++ b/tests/analysis/test_analyze_at_a_size.py @@ -92,7 +92,7 @@ def assert_performance_contract(result: AnalysisResult) -> None: module_target = result.module.resolve_target() throughput = module_target.get_facts(ThroughputFacts) - services = module_target.get_facts(PerformanceServiceFacts, result.topology_level) + services = module_target.get_facts(PerformanceServiceFacts, result.level) scopes = tuple(walk_scopes(build_scopes(result.module, fn))) timed = 0 for expr in collect_exprs(fn.body): @@ -105,7 +105,7 @@ def assert_performance_contract(result: AnalysisResult) -> None: throughput, services, moved=get_metadata(expr, TrafficMetadata), - topology_level=result.topology_level, + level=result.level, ) record = get_metadata(expr, PerformanceMetadata) if not duration: @@ -184,36 +184,37 @@ def _every_number_counts_something(result: AnalysisResult) -> None: """ fn = result.function for expr in (fn, *collect_exprs(fn.body)): - for record in ( - ComputeCostMetadata, - TrafficMetadata, - MemoryMetadata, - RooflineMetadata, - PerformanceMetadata, + for record, rows in ( + (ComputeCostMetadata, ()), + (TrafficMetadata, ()), + (MemoryMetadata, ()), + (RooflineMetadata, ()), + (PerformanceMetadata, ()), ): held = get_metadata(expr, record) if held is None: continue + for field in rows: + value = getattr(held, field) + if field == "flops_logical": + assert value >= 0, f"{describe_expr(expr)}: {field} = {value}" + continue + for name, value in value: + assert value >= 0, f"{describe_expr(expr)}: {field}[{name}] = {value}" if record is ComputeCostMetadata: - for field in ("flops", "service"): - for kind, spread in getattr(held, field).kinds: - for where, value in ( - ("total", spread.total), - *zip(held.topologies, spread.per_unit, strict=False), - ): + for field in ("flops", "other_ops"): + breakdown = getattr(held, field) + for name, spread in breakdown.kinds: + for value in (spread.logical, spread.total, *spread.per_unit): assert value >= 0, ( - f"{describe_expr(expr)}: {field}[{kind}].{where} = {value}" + f"{describe_expr(expr)}: {field}[{name}] = {value}" ) if record is TrafficMetadata: - for field in ("storage", "communication"): - for level, spread in getattr(held, field).kinds: - for where, moved in ( - ("total", spread.total), - *zip(held.topologies, spread.per_unit, strict=False), - ): - assert moved.read >= 0 and moved.write >= 0, ( - f"{describe_expr(expr)}: {field}[{level}].{where} = {moved}" - ) + for field in ("whole", "per_unit"): + for level, moved in getattr(held, field): + assert moved.read >= 0 and moved.write >= 0, ( + f"{describe_expr(expr)}: {field}[{level}] = {moved}" + ) for position, moved in enumerate(held.operands): assert moved.read >= 0 and moved.write >= 0, ( f"{describe_expr(expr)}: operand {position} = {moved}" @@ -232,7 +233,7 @@ def _every_number_counts_something(result: AnalysisResult) -> None: record = get_metadata(expr, LoopFootprintMetadata) if record is None: continue - rows = [(item.buffer, item.memory_level) for item in record.footprints] + rows = [(item.buffer, item.level) for item in record.footprints] assert rows == sorted(rows), describe_expr(expr) assert len(rows) == len(set(rows)), describe_expr(expr) for item in record.footprints: @@ -270,9 +271,7 @@ def test_every_analysis_runs_at_a_stated_size(family: str) -> None: def _predicted_ns(module, dims=None) -> int: """What the four families together say one program takes.""" - result = analyze( - module, module.entry_function(), analysis=FAMILIES, topology_level="cta", dims=dims - ) + result = analyze(module, module.entry_function(), analysis=FAMILIES, level="cta", dims=dims) summary = get_metadata(result.function, PerformanceSummaryMetadata) assert summary is not None return summary.timeline.end_ns - summary.timeline.start_ns diff --git a/tests/analysis/test_analyze_cross_module.py b/tests/analysis/test_analyze_cross_module.py index 0a4a9e676..c9fd8c448 100644 --- a/tests/analysis/test_analyze_cross_module.py +++ b/tests/analysis/test_analyze_cross_module.py @@ -51,7 +51,8 @@ def _flops(records) -> dict[str, int]: for item in records: record = item[0] if isinstance(item, tuple) else item for name, spread in record.flops.kinds: - total[name] = total.get(name, 0) + spread.total + value = spread.total + total[name] = total.get(name, 0) + value return total diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index 7345d588a..4dbf9d206 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -166,7 +166,8 @@ def test_analyze_help_explains_topology_effects_and_assumptions(capsys) -> None: help_text = capsys.readouterr().out for family in ("compute-cost", "memory", "roofline", "performance"): assert family in help_text - assert "every level's per-unit share" in help_text + assert "logical" in help_text and "per-unit share" in help_text + assert "traffic" in help_text assert "global traffic is the device's and counted once" in help_text assert "is an observation, not a bound" in help_text @@ -356,7 +357,7 @@ def test_named_provider_registers_only_decorated_targets_and_replays_them( "from tilefoundry.target import Target, register_target\n" "@dataclass(frozen=True)\n" "class _VendorBase(Target):\n" - " pass\n" + " topology_levels = ('core',)\n" "@register_target\n" "@dataclass(frozen=True)\n" "class VendorOne(_VendorBase):\n" @@ -679,18 +680,6 @@ def test_analyze_binds_an_extent_on_a_root_that_reaches_a_child(tmp_path, capsys assert "leaf(" not in expanded -def _shares_text(per_unit: list, topologies: list) -> str: - """One quantity's per-level shares, as the comment renders a dict.""" - return ",".join(f"{name}:{value}" for name, value in zip(topologies, per_unit)) - - -def _bytes_shares_text(per_unit: list, topologies: list) -> str: - """The same for read/write pairs.""" - return ",".join( - f"{name}:r{moved['read']}/w{moved['write']}" for name, moved in zip(topologies, per_unit) - ) - - def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> None: source = Path(__file__).parents[1] / "fixtures" / "placed" / "moe_mega_kernel.py" selector = f"{source}:MoEMegaKernel" @@ -712,7 +701,7 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> assert cli.main(["analyze", selector, str(operands_path), *flags, "--operands"]) == 0 asked = operands_path.read_text(encoding="utf-8") assert "operands=" not in first - assert "operands=0:r30720/w0,result:r0/w30720" in asked + assert "operands=0:r30720/w0;result:r0/w30720" in asked assert ( cli.main(["analyze", selector, str(operands_json_path), *flags, "--operands", "--json"]) == 0 @@ -757,12 +746,14 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> f"# selection requested={','.join(payload['requested'])} " f"executed={','.join(payload['executed'])}", "# compute-cost " - f"flops=f32:{cost['flops']['f32']['total']}" - f"@{_shares_text(cost['flops']['f32']['per_unit'], cost['topologies'])}", + f"flops=f32:{cost['flops']['f32']['logical']}@logical," + f"{cost['flops']['f32']['total']}@total," + f"{cost['flops']['f32']['per_unit'][0]}@{payload['topology']}", "# traffic " f"traffic=gmem:r{moved['storage']['gmem']['total']['read']}" - f"/w{moved['storage']['gmem']['total']['write']}" - f"@{_bytes_shares_text(moved['storage']['gmem']['per_unit'], moved['topologies'])}", + f"/w{moved['storage']['gmem']['total']['write']}@total," + f"r{moved['storage']['gmem']['per_unit'][0]['read']}" + f"/w{moved['storage']['gmem']['per_unit'][0]['write']}@{payload['topology']}", f"# peak-footprint=gmem:{peak[0]['peak_bytes']}", f"# roofline ideal-ns={bound['ideal_ns']} bound-by={bound['bound_by']}", "# performance root=MoEMegaKernel::experts " @@ -770,13 +761,10 @@ def test_analyze_reports_the_inlined_mega_kernel_from_one_rendering(tmp_path) -> f"waves={summary['waves']}", ] assert payload["totals"]["flops"] == { - kind: value["total"] for kind, value in cost["flops"].items() + name: spread["total"] for name, spread in cost["flops"].items() } assert payload["totals"]["traffic"] == { - level: value["total"] for level, value in moved["storage"].items() - } - assert payload["totals"]["communication"] == { - level: value["total"] for level, value in moved["communication"].items() + name: value["total"] for name, value in moved["storage"].items() } hoisted = {line.split(" = ", 1)[0] for line in lines if " = Mesh((Topology(" in line} diff --git a/tests/codegen/test_codegen_registry.py b/tests/codegen/test_codegen_registry.py deleted file mode 100644 index eca233820..000000000 --- a/tests/codegen/test_codegen_registry.py +++ /dev/null @@ -1,100 +0,0 @@ -"""One codegen registry: who answers, from which side of a call, about what. - -A target is a dimension of the key rather than a word in a registry's name, so -a caller of one target asks the callee's target how a call to it is written -without naming anything of that target's. -""" - -from __future__ import annotations - -import pytest - -from tests.fixtures.tir.square import TirSquare -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.codegen.signature import ( - LAUNCH_ABI, - ProgramIdSignature, - ProgramMetaSignature, - TensorSignature, - UnitSignature, - called_as, -) -from tilefoundry.ir.tir.stmts import Sequential -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.target import CpuTarget, CudaTarget -from tilefoundry.visitor_registry.registries import Role, codegen_registry - -_TARGET = CudaTarget("nvidia.h200_sxm") -_TENSOR = TensorSignature( - name="x", - type=TensorType( - shape=(4, DimVar("S", 1, 8)), - dtype=DType.f32, - layout=None, - storage=StorageKind.GMEM, - ), -) - - -def _ctx() -> CudaCodegenContext: - return CudaCodegenContext(target=_TARGET) - - -def test_a_node_is_written_by_the_target_being_written() -> None: - """Emission keys on the same three things as every other question.""" - assert _ctx().handler_for(Sequential) is codegen_registry.lookup( - (CudaTarget, Role.EMIT, Sequential) - ) - - -def test_a_tensor_declares_as_a_pointer_and_the_axes_its_type_leaves_open() -> None: - """A static extent costs the convention nothing; an open one travels along.""" - assert _ctx().declare(_TENSOR) == ("float* x", "int x_shape_1") - - -def test_what_a_convention_adds_declares_in_the_type_it_states() -> None: - ctx = _ctx() - gpu_id = ProgramIdSignature(name="tilefoundry_gpu_program_id", topology_level="gpu") - assert ctx.declare(LAUNCH_ABI[0]) == ("int grid_x",) - assert ctx.declare(LAUNCH_ABI[-1]) == ("void* stream",) - assert ctx.declare(gpu_id) == ("long long tilefoundry_gpu_program_id",) - assert ctx.declare(ProgramMetaSignature(name="tilefoundry_meta")) == ( - "tilefoundry::ProgramMetaData tilefoundry_meta", - ) - - -def test_the_same_tensor_declares_differently_for_two_targets() -> None: - """The host takes the runtime's tensor, which carries its own extents. - - Both answers come out of one registry and the key says which target gave - it; the host's reads nothing off the context it is handed. - """ - host = codegen_registry.lookup((CpuTarget, Role.CALLEE, TensorSignature)) - assert host(_TENSOR, _ctx()) == ("tvm::ffi::Tensor x",) - - -def test_a_caller_asks_the_callees_target_how_the_call_is_written() -> None: - """The argument list is the callee's answer; the names in it are the caller's.""" - ctx = _ctx() - assert ctx.argument(_TENSOR, _TARGET) == ("x", "x_shape_1") - with pytest.raises(RuntimeError, match=r"nothing registered for \(CpuTarget, Role.CALLER"): - ctx.argument(_TENSOR, CpuTarget()) - - -def test_an_unanswered_question_names_the_whole_key() -> None: - """A miss says which target, which side, and which class went unanswered.""" - with pytest.raises( - RuntimeError, match=r"nothing registered for \(CudaTarget, Role.CALLEE, UnitSignature\)" - ): - _ctx().declare(UnitSignature()) - - -def test_how_a_function_is_called_is_stated_by_its_own_target() -> None: - device = TirSquare.lookup("square_device") - host = TirSquare.lookup("square_host") - assert called_as(device).name == "tilefoundry_square_device_launch" - assert called_as(device).trailing == LAUNCH_ABI - assert called_as(host).name == "tilefoundry_square_host_host" - assert called_as(host).trailing == () diff --git a/tests/codegen/test_device_dispatch.py b/tests/codegen/test_device_dispatch.py index ce924ee8f..a45c52510 100644 --- a/tests/codegen/test_device_dispatch.py +++ b/tests/codegen/test_device_dispatch.py @@ -9,30 +9,40 @@ from __future__ import annotations from tests.fixtures.tir.square import TirSquare +from tilefoundry.codegen.context_builder import build_codegen_context from tilefoundry.codegen.cpu.context import CpuCodegenContext from tilefoundry.codegen.cpu.module import emit_host_module from tilefoundry.codegen.cuda.context import CudaCodegenContext from tilefoundry.codegen.cuda.module import emit_cuda_module -from tilefoundry.codegen.signature import symbol_table -from tilefoundry.codegen.topology import launch_geometry from tilefoundry.target import CudaTarget _CUDA = CudaTarget("nvidia.h200_sxm") +def _view(target): + root = build_codegen_context(TirSquare) + group = next(group for group in root.groups if group.target == target) + return root, group, root.for_group(group) + + def _device_source() -> str: + _root, group, view = _view(_CUDA) ctx = CudaCodegenContext( - symbols=symbol_table(TirSquare, _CUDA), + symbols=view.symbols, target=_CUDA, - launches=launch_geometry(TirSquare), + launches=view.launches, ) - kernels = tuple(fn for fn in TirSquare.functions if fn.target == _CUDA) - return emit_cuda_module(TirSquare, kernels, _CUDA, ctx).source + return emit_cuda_module(TirSquare, group.functions, _CUDA, ctx).source def _host_source() -> str: entry = TirSquare.entry_function() - ctx = CpuCodegenContext(symbols=symbol_table(TirSquare, _CUDA), target=entry.target) + _root, group, view = _view(entry.target) + ctx = CpuCodegenContext( + symbols=view.symbols, + target=entry.target, + resolved_launches=view.resolved_launches, + ) return emit_host_module(TirSquare, (entry,), entry.target, ctx).source diff --git a/tests/codegen/test_host_multi_launch.py b/tests/codegen/test_host_multi_launch.py index 7f2fcf9d3..cb01b542e 100644 --- a/tests/codegen/test_host_multi_launch.py +++ b/tests/codegen/test_host_multi_launch.py @@ -7,9 +7,9 @@ import tilefoundry from tilefoundry import module, prim_func +from tilefoundry.codegen.context_builder import build_codegen_context from tilefoundry.codegen.cpu.context import CpuCodegenContext from tilefoundry.codegen.cpu.module import emit_host_module -from tilefoundry.codegen.signature import symbol_table from tilefoundry.dsl import T, Tensor from tilefoundry.ir.core import Constant, Var from tilefoundry.ir.core.module import Module @@ -129,8 +129,13 @@ def _module(devices, launches, params) -> tuple[Module, PrimFunction]: def _source(devices, launches, params) -> str: mod, host = _module(devices, launches, params) + root = build_codegen_context(mod) + group = next(group for group in root.groups if host in group.functions) + view = root.for_group(group) ctx = CpuCodegenContext( - symbols=symbol_table(mod, CudaTarget("nvidia.h200_sxm")), target=host.target + symbols=view.symbols, + target=host.target, + resolved_launches=view.resolved_launches, ) return emit_host_module(mod, (host,), host.target, ctx).source @@ -144,8 +149,8 @@ def test_two_launches_keep_order_and_each_block_size() -> None: (a, b), ) assert source.index("tilefoundry_first_launch") < source.index("tilefoundry_second_launch") - assert "1, 32, 1, 1" in source - assert "1, 64, 1, 1" in source + assert "1, 32, 0" in source + assert "1, 64, 0" in source @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") diff --git a/tests/codegen/test_linkable.py b/tests/codegen/test_linkable.py index 41aa72af7..273a4cb54 100644 --- a/tests/codegen/test_linkable.py +++ b/tests/codegen/test_linkable.py @@ -3,11 +3,10 @@ from __future__ import annotations from tests.fixtures.tir.square import TirSquare +from tilefoundry.codegen.context_builder import build_codegen_context from tilefoundry.codegen.cpu.context import CpuCodegenContext from tilefoundry.codegen.cpu.module import emit_host_module from tilefoundry.codegen.linkable import LinkableFunction, LinkableModule -from tilefoundry.codegen.signature import symbol_table -from tilefoundry.target import CudaTarget def _linkable_function(name: str) -> LinkableFunction: @@ -27,8 +26,13 @@ def _param_types(parameter_list: str) -> list[str]: def _host_unit() -> LinkableModule: """The host translation unit of a module whose entry launches a shape-dispatched kernel.""" entry = TirSquare.entry_function() + root = build_codegen_context(TirSquare) + group = next(group for group in root.groups if entry in group.functions) + view = root.for_group(group) ctx = CpuCodegenContext( - symbols=symbol_table(TirSquare, CudaTarget("nvidia.h200_sxm")), target=entry.target + symbols=view.symbols, + target=entry.target, + resolved_launches=view.resolved_launches, ) return emit_host_module(TirSquare, (entry,), entry.target, ctx) diff --git a/tests/codegen/test_symbol_table.py b/tests/codegen/test_symbol_table.py deleted file mode 100644 index 397923b4f..000000000 --- a/tests/codegen/test_symbol_table.py +++ /dev/null @@ -1,98 +0,0 @@ -"""The table every emitter reads before it writes a line. - -One row per function a caller can name, settled once for a whole compile, so -a host entry and the shim it calls cannot describe the same call differently. -""" - -from __future__ import annotations - -import pytest - -from tests.fixtures.placed.gpu_placed_rows import GPUS, GpuPlacedRows -from tests.fixtures.tir.square import TirSquare -from tilefoundry.codegen.cuda.context import CudaCodegenContext -from tilefoundry.codegen.signature import ( - LAUNCH_ABI, - ProgramIdSignature, - TensorSignature, - program_id_params, - symbol_table, -) -from tilefoundry.ir.core.module import module_functions -from tilefoundry.ir.types import DType, TensorType -from tilefoundry.ir.types.dim import DimVar -from tilefoundry.ir.types.storage import StorageKind -from tilefoundry.target import CudaTarget - -_PLACED = CudaTarget("nvidia.h200_sxm", device_count=GPUS) -_GPU_ID = ProgramIdSignature(name="tilefoundry_gpu_program_id", topology_level="gpu") - - -def test_the_ids_a_call_carries_are_the_levels_the_target_states() -> None: - assert program_id_params(GpuPlacedRows, _PLACED) == (_GPU_ID,) - - -def test_a_program_that_names_no_stated_level_carries_no_ids() -> None: - """The same target, a program that never names ``gpu``: nothing to tell it.""" - assert program_id_params(TirSquare, _PLACED) == () - - -def test_a_device_function_is_reached_through_its_shim() -> None: - table = symbol_table(GpuPlacedRows, _PLACED) - shim = table[id(GpuPlacedRows.lookup("copy_rows_device"))] - assert shim.name == "tilefoundry_copy_rows_device_launch" - assert shim.leading == (_GPU_ID,) - assert shim.trailing == LAUNCH_ABI - assert [p.name for p in shim.params] == ["a", "out"] - - -def test_a_host_entry_is_reached_directly_and_told_what_it_places() -> None: - """The id enters the program here: the process is the only thing that knows it.""" - table = symbol_table(GpuPlacedRows, _PLACED) - entry = table[id(GpuPlacedRows.lookup("copy_rows_host"))] - assert entry.name == "tilefoundry_copy_rows_host_host" - assert entry.leading == (_GPU_ID,) - assert entry.trailing == () - - -def test_the_kernels_own_convention_is_not_in_the_table() -> None: - """Only the shim is callable, so only the shim is written down.""" - table = symbol_table(GpuPlacedRows, _PLACED) - assert not any(signature.name.endswith("_kernel") for signature in table.values()) - - -def test_a_prototype_gets_one_row_and_its_variants_none() -> None: - """Nothing calls a variant: the prototype is the one symbol, and it branches.""" - table = symbol_table(TirSquare, _PLACED) - assert sorted(signature.name for signature in table.values()) == [ - "tilefoundry_square_device_launch", - "tilefoundry_square_host_host", - ] - - -def test_a_call_site_looks_its_callee_up_by_identity() -> None: - """Keyed by ``id(fn)``: the call site holds the callee, not its module.""" - assert set(symbol_table(TirSquare, _PLACED)) == { - id(fn) for fn in module_functions(TirSquare) - } - - -def test_the_context_answers_from_the_table_it_was_given() -> None: - device = GpuPlacedRows.lookup("copy_rows_device") - ctx = CudaCodegenContext(symbols=symbol_table(GpuPlacedRows, _PLACED), target=_PLACED) - assert ctx.signature_of(device).name == "tilefoundry_copy_rows_device_launch" - with pytest.raises(KeyError, match="symbol table"): - ctx.signature_of(TirSquare.lookup("square_host")) - - -def test_the_dynamic_axes_are_the_ones_the_type_leaves_open() -> None: - signature = TensorSignature( - name="x", - type=TensorType( - shape=(4, DimVar("S", 1, 8)), - dtype=DType.f32, - layout=None, - storage=StorageKind.GMEM, - ), - ) - assert signature.dynamic_axes == (1,) diff --git a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt index f3f70372a..b6ec3c9eb 100644 --- a/tests/fixtures/inspection/type_printer_sugar.analyzed.txt +++ b/tests/fixtures/inspection/type_printer_sugar.analyzed.txt @@ -19,21 +19,21 @@ def composed_mesh_pipeline( ), "rmem"] ): with Mesh((Topology("cta", 4),), Layout((4,), (1,)), names=('tile',)) as cta: - v0 = reshard(x, layout=(4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r2048/w0@cta:r512/w0,thread:r64/w0,rmem:r0/w2048@cta:r0/w512,thread:r0/w64 - v1 = unary(v0, kind="square") # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost flops=f32:512@cta:128,thread:16; traffic traffic=rmem:r2048/w2048@cta:r512/w512,thread:r64/w64 - v2 = reshard(v1, layout=(4 @ cta.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost; traffic traffic=rmem:r2048/w0@cta:r512/w0,thread:r64/w0,smem:r0/w2048@cta:r0/w512,thread:r0/w64 - v3 = cast(v2, dtype="bf16") # Tensor[(8, 4, 16), "bf16", ((4 @ cta.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost flops=bf16:512@cta:128,thread:128; traffic traffic=smem:r2048/w1024@cta:r512/w256,thread:r512/w256 - v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ cta.tile, 2, 16, 4), (128, 64, 1, 16)), "smem"]; compute-cost; traffic traffic=smem:r1024/w1024@cta:r256/w256,thread:r256/w256 - v5 = reshard(v4, layout=((8, 16, 4), (64, 4, 1), {cta.tile @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {cta.tile @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@cta:r0/w256,thread:r0/w256,smem:r1024/w0@cta:r256/w0,thread:r256/w0 - v6 = reshard(seed, layout=((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), storage=rmem) # Tensor[(16,), "f32", ((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r64/w0@cta:r64/w0,thread:r8/w0,rmem:r0/w64@cta:r0/w64,thread:r0/w8 + v0 = reshard(x, layout=(4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), storage=rmem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r2048/w0@total,r512/w0@cta,r64/w0@thread;rmem:r0/w2048@total,r0/w512@cta,r0/w64@thread + v1 = unary(v0, kind="square") # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 2 @ mesh.warp, 2, 4 @ mesh.lane, 4), (0, 8, 0, 4, 0, 1)), "rmem"]; compute-cost flops=f32:512@logical,512@total,128@cta,16@thread; traffic traffic=rmem:r2048/w2048@total,r512/w512@cta,r64/w64@thread + v2 = reshard(v1, layout=(4 @ cta.tile, 2, 4, 16), storage=smem) # Tensor[(8, 4, 16), "f32", ((4 @ cta.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost; traffic traffic=rmem:r2048/w0@total,r512/w0@cta,r64/w0@thread;smem:r0/w2048@total,r0/w512@cta,r0/w64@thread + v3 = cast(v2, dtype="bf16") # Tensor[(8, 4, 16), "bf16", ((4 @ cta.tile, 2, 4, 16), (128, 64, 16, 1)), "smem"]; compute-cost flops=bf16:512@logical,512@total,128@cta,128@thread; traffic traffic=smem:r2048/w1024@total,r512/w256@cta,r512/w256@thread + v4 = transpose(v3, perm=(0, 2, 1)) # Tensor[(8, 16, 4), "bf16", ((4 @ cta.tile, 2, 16, 4), (128, 64, 1, 16)), "smem"]; compute-cost; traffic traffic=smem:r1024/w1024@total,r256/w256@cta,r256/w256@thread + v5 = reshard(v4, layout=((8, 16, 4), (64, 4, 1), {cta.tile @ B()}), storage=gmem) # Tensor[(8, 16, 4), "bf16", ((8, 16, 4), (64, 4, 1), {cta.tile @ B()})]; compute-cost; traffic traffic=gmem:r0/w1024@total,r0/w256@cta,r0/w256@thread;smem:r1024/w0@total,r256/w0@cta,r256/w0@thread + v6 = reshard(seed, layout=((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), storage=rmem) # Tensor[(16,), "f32", ((2 @ mesh.warp, 4 @ mesh.lane, 2), (8, 2, 1)), "rmem"]; compute-cost; traffic traffic=gmem:r64/w0@total,r64/w0@cta,r8/w0@thread;rmem:r0/w64@total,r0/w64@cta,r0/w8@thread folded = reshard(acc, layout=(2 @ mesh.warp, 4, 4 @ mesh.lane, 4), storage=rmem) # Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost; traffic summed = reshard(mixed, layout=((8, 16), {mesh.warp @ B()}), storage=rmem) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]; compute-cost; traffic - for _ in range(3): # loop-footprint footprints=folded@rmem:8192/196608/24576,summed@rmem:131072/393216/393216 status=complete - v10 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]; compute-cost flops=f32:128@cta:128,thread:128; traffic traffic=rmem:r1024/w512@cta:r1024/w512,thread:r1024/w512 - v11 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:128@cta:128,thread:16; traffic traffic=rmem:r512/w512@cta:r512/w512,thread:r64/w64 + for _ in range(3): # loop-footprint footprints=folded@rmem:8192/196608/24576;summed@rmem:131072/393216/393216 status=complete + v10 = add(summed, summed) # Tensor[(8, 16), "f32", ((8, 16), (16, 1), {mesh.warp @ B()}), "rmem"]; compute-cost flops=f32:128@logical,128@total,128@cta,128@thread; traffic traffic=rmem:r1024/w512@total,r1024/w512@cta,r1024/w512@thread + v11 = unary(folded, kind="square") # Tensor[(8, 16), "f32", ((2 @ mesh.warp, 4, 4 @ mesh.lane, 4), (64, 16, 4, 1)), "rmem"]; compute-cost flops=f32:128@logical,128@total,128@cta,16@thread; traffic traffic=rmem:r512/w512@total,r512/w512@cta,r64/w64@thread folded = v11 summed = v10 - v13 = reshard(folded, layout=Layout((8, 16), (16, 1)), storage=gmem) # Tensor[(8, 16), "f32", Layout((8, 16), (16, 1))]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w64,rmem:r512/w0@cta:r512/w0,thread:r64/w0 - v15 = reshard(summed, layout=Layout((8, 16), (16, 1)), storage=gmem) # Tensor[(8, 16), "f32", Layout((8, 16), (16, 1))]; compute-cost; traffic traffic=gmem:r0/w512@cta:r0/w512,thread:r0/w512,rmem:r512/w0@cta:r512/w0,thread:r512/w0 - v17 = reshard(v6, layout=Layout((16,), (1,)), storage=gmem) # Tensor[(16,), "f32", Layout((16,), (1,))]; compute-cost; traffic traffic=gmem:r0/w64@cta:r0/w64,thread:r0/w8,rmem:r64/w0@cta:r64/w0,thread:r8/w0 + v13 = reshard(folded, layout=Layout((8, 16), (16, 1)), storage=gmem) # Tensor[(8, 16), "f32", Layout((8, 16), (16, 1))]; compute-cost; traffic traffic=gmem:r0/w512@total,r0/w512@cta,r0/w64@thread;rmem:r512/w0@total,r512/w0@cta,r64/w0@thread + v15 = reshard(summed, layout=Layout((8, 16), (16, 1)), storage=gmem) # Tensor[(8, 16), "f32", Layout((8, 16), (16, 1))]; compute-cost; traffic traffic=gmem:r0/w512@total,r0/w512@cta,r0/w512@thread;rmem:r512/w0@total,r512/w0@cta,r512/w0@thread + v17 = reshard(v6, layout=Layout((16,), (1,)), storage=gmem) # Tensor[(16,), "f32", Layout((16,), (1,))]; compute-cost; traffic traffic=gmem:r0/w64@total,r0/w64@cta,r0/w8@thread;rmem:r64/w0@total,r64/w0@cta,r8/w0@thread return (v5, v13, v15, v17) diff --git a/tests/fixtures/tir/async_sync.py b/tests/fixtures/tir/async_sync.py index 682deefe4..73acbf7f0 100644 --- a/tests/fixtures/tir/async_sync.py +++ b/tests/fixtures/tir/async_sync.py @@ -6,7 +6,7 @@ from tilefoundry.target import CpuTarget, CudaTarget -@module(entry="async_stage_host", target=CudaTarget("nvidia.h200_sxm")) +@module(entry="async_stage_host", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("thread", 128),)) class AsyncStage: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def async_stage_device(a: Tensor[(128, 4), "f32"], b: Tensor[(128, 4), "f32"]): diff --git a/tests/fixtures/tir/square.py b/tests/fixtures/tir/square.py index 038e795ed..2d5737358 100644 --- a/tests/fixtures/tir/square.py +++ b/tests/fixtures/tir/square.py @@ -10,7 +10,7 @@ _S = DimVar("S", 1, 256) -@module(entry="square_host", target=CudaTarget("nvidia.h200_sxm")) +@module(entry="square_host", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("thread", 128),)) class TirSquare: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def square_device(x: Tensor[(_S,), "f32"]): diff --git a/tests/fixtures/tir/sync.py b/tests/fixtures/tir/sync.py index 7d461e54c..43b95be3c 100644 --- a/tests/fixtures/tir/sync.py +++ b/tests/fixtures/tir/sync.py @@ -7,7 +7,7 @@ from tilefoundry.target import CpuTarget, CudaTarget -@module(entry="sync_square_host", target=CudaTarget("nvidia.h200_sxm")) +@module(entry="sync_square_host", target=CudaTarget("nvidia.h200_sxm"), topologies=(Topology("thread", 128),)) class SyncSquare: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def sync_square_device(a: Tensor[(4, 32), "f32"]): diff --git a/tests/installed/models/smoke_qwen3_1_7b.py b/tests/installed/models/smoke_qwen3_1_7b.py index bc88c2941..bdc0fd5e0 100644 --- a/tests/installed/models/smoke_qwen3_1_7b.py +++ b/tests/installed/models/smoke_qwen3_1_7b.py @@ -230,10 +230,10 @@ def test_the_mlp_costs_its_three_matrices(tf, shipped_source) -> None: authored = _reported(tf, source, "mlp", None) placed = _reported(tf, source, "placed_mlp", None) authored_flops = authored["totals"]["flops"][DT] - stated = placed["function_records"]["compute-cost"]["flops"][DT] + stated = placed["function_records"]["compute-cost"]["flops"] topologies = placed["function_records"]["compute-cost"]["topologies"] - placed_flops = stated["total"] - per_unit = dict(zip(topologies, stated["per_unit"])) + placed_flops = stated[DT]["total"] + per_unit = dict(zip(topologies, stated[DT]["per_unit"])) assert placed_flops == per_unit["cta"] * 128 _holds(authored_flops, _mlp_matmul_flops(), "mlp") diff --git a/tests/ops/ir/cost_utils.py b/tests/ops/ir/cost_utils.py index 73f469357..ca3f7a7cd 100644 --- a/tests/ops/ir/cost_utils.py +++ b/tests/ops/ir/cost_utils.py @@ -20,7 +20,7 @@ class CostCase: op: object inputs: tuple[TensorType, ...] flops: Mapping[DType, int] = field(default_factory=dict) - service: Mapping[str, int] = field(default_factory=dict) + ops: Mapping[str, int] = field(default_factory=dict) traffic: tuple[TrafficBytes, ...] = () topology_level: str | None = None topologies: tuple[Topology, ...] = () @@ -44,5 +44,5 @@ def run_cost_case(case: CostCase) -> None: cost = CostEvaluator().visit_Call(call, ctx) assert cost.flops == case.flops - assert cost.service == case.service + assert cost.ops == case.ops assert cost.traffic == case.traffic diff --git a/tests/ops/ir/test_binary.py b/tests/ops/ir/test_binary.py index 7ca0c1984..6c19623b7 100644 --- a/tests/ops/ir/test_binary.py +++ b/tests/ops/ir/test_binary.py @@ -173,7 +173,7 @@ def test_lower_rank_split_right_aligns(): "lt_over_reals_is_a_predicate", Binary(kind=BinaryKind.LT), (_REAL, _REAL), - service={"predicate": 8}, + ops={"predicate": 8}, traffic=(TrafficBytes(read=32), TrafficBytes(read=32), TrafficBytes(write=1)), ), CostCase( diff --git a/tests/ops/ir/test_unary.py b/tests/ops/ir/test_unary.py index 14756c8ec..8013e8032 100644 --- a/tests/ops/ir/test_unary.py +++ b/tests/ops/ir/test_unary.py @@ -101,7 +101,7 @@ def test_a_unary_evaluates_to_its_oracle(name, op, drawn, oracle): "exp_is_special_function_work", Unary(kind=UnaryKind.EXP), (_REAL,), - service={"special": 8}, + ops={"special": 8}, traffic=(TrafficBytes(read=32), TrafficBytes(write=32)), ), CostCase( diff --git a/tests/ops/ir/test_where.py b/tests/ops/ir/test_where.py index baf2045ac..facda266f 100644 --- a/tests/ops/ir/test_where.py +++ b/tests/ops/ir/test_where.py @@ -119,7 +119,7 @@ def test_where_cost_counts_selection_and_materialization(): make_tensor_type((2, 3, 4), DType.f32), make_tensor_type((2, 3, 4), DType.f32), ), - service={"select": 24}, + ops={"select": 24}, traffic=( TrafficBytes(read=3), TrafficBytes(read=96), diff --git a/tests/ops/tir/cuda/test_mma.py b/tests/ops/tir/cuda/test_mma.py index 08293a9b6..77d0b14d8 100644 --- a/tests/ops/tir/cuda/test_mma.py +++ b/tests/ops/tir/cuda/test_mma.py @@ -35,7 +35,7 @@ def test_handwritten_mma_matches_torch() -> None: torch.testing.assert_close(out, torch.matmul(a.float(), b.float()), rtol=2e-2, atol=2e-2) -@module(entry="tile_host", target=_CUDA) +@module(entry="tile_host", target=_CUDA, topologies=(Topology("thread", 32),)) class MmaTile: """A 16x16 by 16x8 product done as one atom, looped by the tile's shape.""" @@ -68,22 +68,30 @@ def tile_device( ), ) a_tile = T.alloc_tensor( - Tensor[(16, 16), 'bf16', ShardLayout( + Tensor[ + (16, 16), + "bf16", + ShardLayout( layout=Layout(shape=(16, 16), strides=(16, 1)), attrs=(Broadcast(), Broadcast()), mesh=m, - ), 'smem'] + ), + "smem", + ] ) b_tile = T.alloc_tensor( - Tensor[(8, 16), 'bf16', ShardLayout( + Tensor[ + (8, 16), + "bf16", + ShardLayout( layout=Layout(shape=(8, 16), strides=(1, 8)), attrs=(Broadcast(), Broadcast()), mesh=m, - ), 'smem'] - ) - acc = T.alloc_tensor( - Tensor[(16, 8), 'f32', atom.C, 'rmem'] + ), + "smem", + ] ) + acc = T.alloc_tensor(Tensor[(16, 8), "f32", atom.C, "rmem"]) T.copy(a_view, a_tile) T.copy(b_view, b_tile) T.fill(acc, 0.0) diff --git a/tests/ops/tir/cuda/test_tma.py b/tests/ops/tir/cuda/test_tma.py index fa77cab3e..e4a558ab1 100644 --- a/tests/ops/tir/cuda/test_tma.py +++ b/tests/ops/tir/cuda/test_tma.py @@ -73,7 +73,10 @@ def test_refuses_a_shape_change() -> None: verify_prim_function(_pf(_ty(8), _ty(4, storage="smem"))) -@pytest.mark.parametrize(("n", "dtype"), [(5, DType.from_name("f32")), (1, DType.from_name("f32")), (7, DType.from_name("bf16"))]) +@pytest.mark.parametrize( + ("n", "dtype"), + [(5, DType.from_name("f32")), (1, DType.from_name("f32")), (7, DType.from_name("bf16"))], +) def test_admits_a_transfer_off_the_sixteen_byte_grain(n, dtype) -> None: """The grain belongs to one instruction, and the op does not name one. @@ -82,7 +85,11 @@ def test_admits_a_transfer_off_the_sixteen_byte_grain(n, dtype) -> None: verify_prim_function(_pf(_ty(n, dtype), _ty(n, dtype, storage="smem"))) -@module(entry="tma_tiers_host", target=CudaTarget("nvidia.h200_sxm")) +@module( + entry="tma_tiers_host", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("thread", 128),), +) class TmaTiers: """Both staging tiers in one device function, one operand pair each.""" @@ -101,15 +108,18 @@ def tma_tiers_device( ), ) bulk_stage = T.alloc_tensor( - Tensor[(256,), 'f32', ShardLayout( + Tensor[ + (256,), + "f32", + ShardLayout( layout=Layout(shape=(256,), strides=(1,)), attrs=(Broadcast(),), mesh=m, - ), 'smem'] - ) - bulk_bar = T.alloc_tensor( - Tensor[(1,), 'i64', None, 'smem'] + ), + "smem", + ] ) + bulk_bar = T.alloc_tensor(Tensor[(1,), "i64", None, "smem"]) T.mbarrier_init(bulk_bar, arrive_count=1) T.sync(m) T.tma_copy(bulk_view, bulk_stage, bulk_bar) @@ -123,15 +133,18 @@ def tma_tiers_device( ), ) odd_stage = T.alloc_tensor( - Tensor[(5,), 'f32', ShardLayout( + Tensor[ + (5,), + "f32", + ShardLayout( layout=Layout(shape=(5,), strides=(1,)), attrs=(Broadcast(),), mesh=mo, - ), 'smem'] - ) - odd_bar = T.alloc_tensor( - Tensor[(1,), 'i64', None, 'smem'] + ), + "smem", + ] ) + odd_bar = T.alloc_tensor(Tensor[(1,), "i64", None, "smem"]) T.mbarrier_init(odd_bar, arrive_count=1) T.sync(mo) T.tma_copy(odd_view, odd_stage, odd_bar) diff --git a/tests/ops/tir/test_copy.py b/tests/ops/tir/test_copy.py index 02583aa2d..3095c4a75 100644 --- a/tests/ops/tir/test_copy.py +++ b/tests/ops/tir/test_copy.py @@ -19,7 +19,7 @@ _CUDA = CudaTarget("nvidia.h200_sxm") -@module(entry="copy_storage_host", target=_CUDA) +@module(entry="copy_storage_host", target=_CUDA, topologies=(Topology("thread", 128),)) class CopyStorage: """Five operand pairs in one device function, one per thing that can differ.""" @@ -39,44 +39,32 @@ def copy_storage_device( with Mesh((Topology("thread", 128),), Layout(shape=(128,), strides=(1,)), ("t",)) as m: smem_src = T.tensor_view(a_smem, layout=split_rows(m)) smem_dst = T.tensor_view(b_smem, layout=split_rows(m)) - smem_tile = T.alloc_tensor( - Tensor[(128, 4), 'f32', split_rows(m), 'smem'] - ) + smem_tile = T.alloc_tensor(Tensor[(128, 4), "f32", split_rows(m), "smem"]) T.copy(smem_src, smem_tile) T.sync(m) T.copy(smem_tile, smem_dst) rmem_src = T.tensor_view(a_rmem, layout=split_rows(m)) rmem_dst = T.tensor_view(b_rmem, layout=split_rows(m)) - rmem_tile = T.alloc_tensor( - Tensor[(128, 4), 'f32', split_rows(m), 'smem'] - ) - rmem_frag = T.alloc_tensor( - Tensor[(128, 4), 'f32', split_rows(m), 'rmem'] - ) + rmem_tile = T.alloc_tensor(Tensor[(128, 4), "f32", split_rows(m), "smem"]) + rmem_frag = T.alloc_tensor(Tensor[(128, 4), "f32", split_rows(m), "rmem"]) T.copy(rmem_src, rmem_tile) T.sync(m) T.copy(rmem_tile, rmem_frag) T.copy(rmem_frag, rmem_dst) wide_src = T.tensor_view(a_wide, layout=split_rows(m)) wide_dst = T.tensor_view(b_wide, layout=split_rows(m)) - wide_frag = T.alloc_tensor( - Tensor[(128, 4), 'f32', split_rows(m), 'rmem'] - ) + wide_frag = T.alloc_tensor(Tensor[(128, 4), "f32", split_rows(m), "rmem"]) T.copy(wide_src, wide_frag) T.copy(wide_frag, wide_dst) narrow_src = T.tensor_view(a_narrow, layout=split_pairs(m)) narrow_dst = T.tensor_view(b_narrow, layout=split_pairs(m)) - narrow_frag = T.alloc_tensor( - Tensor[(128, 2), 'f32', split_pairs(m), 'rmem'] - ) + narrow_frag = T.alloc_tensor(Tensor[(128, 2), "f32", split_pairs(m), "rmem"]) T.copy(narrow_src, narrow_frag) T.copy(narrow_frag, narrow_dst) with Mesh((Topology("thread", 32),), Layout(shape=(32,), strides=(1,)), ("t",)) as mb: bcast_src = T.tensor_view(a_bcast, layout=split_short_rows(mb)) bcast_dst = T.tensor_view(b_bcast, layout=split_short_rows(mb)) - bcast_frag = T.alloc_tensor( - Tensor[(4,), 'f32', broadcast_run(mb), 'rmem'] - ) + bcast_frag = T.alloc_tensor(Tensor[(4,), "f32", broadcast_run(mb), "rmem"]) T.copy(bcast_src, bcast_frag) T.copy(bcast_frag, bcast_dst) diff --git a/tests/ops/tir/test_sync.py b/tests/ops/tir/test_sync.py index 03aba153b..77fffb7fd 100644 --- a/tests/ops/tir/test_sync.py +++ b/tests/ops/tir/test_sync.py @@ -18,7 +18,7 @@ _CUDA = CudaTarget("nvidia.h200_sxm") -@module(entry="slice_host", target=_CUDA) +@module(entry="slice_host", target=_CUDA, topologies=(Topology("thread", 128),)) class SyncSlices: @prim_func(target=_CUDA) def one_warp(a: Tensor[(1, 32), "f32"], o: Tensor[(1, 32), "f32"]): @@ -55,7 +55,11 @@ def slice_host( launch(two_warps, a1, o1, grid=(1, 1, 1), block=(128, 1, 1)) # noqa: F821 -@module(entry="grid_sync_host", target=_CUDA) +@module( + entry="grid_sync_host", + target=_CUDA, + topologies=(Topology("cta", 4), Topology("thread", 128)), +) class GridSync: @prim_func(target=CudaTarget("nvidia.h200_sxm")) def grid_sync_device(a: Tensor[(128,), "f32"]): diff --git a/tests/passes/test_host_entry.py b/tests/passes/test_host_entry.py index a3b9a7fab..ed4452b9f 100644 --- a/tests/passes/test_host_entry.py +++ b/tests/passes/test_host_entry.py @@ -67,29 +67,6 @@ def large(x: Tensor[(_S,), "f32"]): T.sync(thread) -@module(entry="square", topologies=(Topology("thread", 128),)) -class _DisagreeingVariants: - """The same shape, except the two variants mesh over different thread counts.""" - - @prim_func(target=_CUDA) - def square(x: Tensor[(_S,), "f32"]): - pass - - @square.specialize(DimVarRangePat("S", 1, 127)) - def small(x: Tensor[(_S,), "f32"]): - with Mesh((Topology("thread", 64),), Layout((64,), (1,))) as thread: - view = T.tensor_view(x, layout=_rows(64)) - T.copy(view, view) - T.sync(thread) - - @square.specialize(DimVarRangePat("S", 128, 255)) - def large(x: Tensor[(_S,), "f32"]): - with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread: - view = T.tensor_view(x, layout=_rows(128)) - T.copy(view, view) - T.sync(thread) - - def _launch_of(mod): """The one ``Launch`` the synthesized entry's body is made of.""" (statement,) = mod.entry_function().body.body @@ -128,10 +105,5 @@ def test_a_lone_kernel_reads_its_geometry_off_its_own_body() -> None: assert _geometry(_launch_of(insert_default_host_entry(_OneKernel))) == ((1, 1, 1), (128, 1, 1)) -def test_variants_that_want_different_geometries_are_refused() -> None: - with pytest.raises(ValueError, match="different geometries"): - insert_default_host_entry(_DisagreeingVariants) - - def test_a_module_that_already_has_a_host_entry_is_left_alone() -> None: assert insert_default_host_entry(TirSquare) is TirSquare diff --git a/tests/target/test_target.py b/tests/target/test_target.py index 9736e566d..34ff260f7 100644 --- a/tests/target/test_target.py +++ b/tests/target/test_target.py @@ -20,7 +20,6 @@ from tests.installed.smoke_target.vendor_npu import VendorNpuTarget from tilefoundry import CompilerOptions, DType, build, jit, module from tilefoundry.analysis import AnalysisError, analyze -from tilefoundry.codegen.registry import group_functions_by_target from tilefoundry.dsl import DimVar from tilefoundry.ir.core.module import Module from tilefoundry.ir.tir.prim_function import PrimFunction @@ -126,9 +125,7 @@ def test_document_free_target_projects_facts_and_inherits_standard_analysis() -> target = VendorNpuTarget() assert target.get_analyzer("roofline").selector == "roofline" - assert target.get_facts(TopologyLimitFacts, "core") == TopologyLimitFacts( - "core", 256 - ) + assert target.get_facts(TopologyLimitFacts, "core") == TopologyLimitFacts("core", 256) memory = target.get_facts(MemoryHierarchyFacts) assert memory.explicit("gmem").capacity_bytes == 64_000_000_000 assert target.get_facts(ThroughputFacts).peak_for(DType.f32) == 2_000_000_000_000_000 @@ -152,8 +149,7 @@ def test_cuda_projects_one_ctas_share_from_device_rates() -> None: assert services.unit == "cta" assert dict(services.unit_flops) == { - dtype: rate // target.device.sm_count - for dtype, rate in device.peak_flops_per_second + dtype: rate // target.device.sm_count for dtype, rate in device.peak_flops_per_second } assert dict(services.unit_bandwidth) == { device.bandwidth_level: ( @@ -289,36 +285,6 @@ def test_program_topologies_use_target_resource_facts() -> None: assert len(str(error.value)) < 120 -def test_group_functions_by_target_fact_matching() -> None: - """CUDA functions must agree on Target facts before grouping. - - CUDA functions must agree on Target facts before grouping; CPU - functions are exempt from the CUDA fact-matching. - """ - body = Sequential(body=()) - first = PrimFunction(name="first", params=(), body=body, target=CudaTarget("nvidia.h200_sxm")) - second = PrimFunction( - name="second", - params=(), - body=body, - target=CudaTarget( - "nvidia.h200_sxm", - architecture=replace(CudaTarget.hardware.resolve(SM90_ID).value, name="sm_90_alt"), - ), - ) - with pytest.raises(ValueError, match="mixes unequal device Targets") as error: - group_functions_by_target(Module(name="mixed", functions=(first, second), entry="first")) - assert "CudaTarget (cuda)" in str(error.value) - assert "architecture=" not in str(error.value) - assert "device=" not in str(error.value) - assert len(str(error.value)) < 300 - - host = PrimFunction(name="host", params=(), body=body, target=CpuTarget()) - groups = group_functions_by_target(Module(name="mixed", functions=(first, host), entry="host")) - assert tuple(fn.name for fn in groups[first.target]) == ("first",) - assert tuple(fn.name for fn in groups[host.target]) == ("host",) - - def test_build_rejects_unequal_device_targets_before_emitting() -> None: emitted: list[object] = [] @@ -349,12 +315,7 @@ def get_code_generator(self) -> CodeGenerator: with pytest.raises(ValueError) as error: build(mixed) - assert str(error.value) == ( - "tilefoundry: module 'mixed' mixes unequal device Targets: " - "FirstDeviceTarget (tests.target.first_device) (function 'first') vs " - "SecondDeviceTarget (tests.target.second_device) (function 'second'); " - "multiple device translation units are not supported" - ) + assert str(error.value) == "codegen: expected one device Target, found 2" assert emitted == []