Skip to content

Commit 7d23c61

Browse files
tanner-andrulistanner-andrulis
committed
Typo fixes
1 parent 961771e commit 7d23c61

2 files changed

Lines changed: 94 additions & 62 deletions

File tree

‎examples/arches/orojenesis.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ arch:
2424
- !Memory
2525
name: GlobalBuffer
2626
component_class: Dummy
27-
size: 1e12
27+
size: 1e15
2828
tensors: {keep: All}
2929

3030
- !Compute

‎notebooks/tutorials/memory_size_vs_access_tradeoff.ipynb‎

Lines changed: 93 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,8 @@
1313
"\n",
1414
"This analysis is called Orojenesis, and it is introduced in \"Mind the Gap: Attainable\n",
1515
"Data Movement and Operational Intensity Bounds for Tensor Algorithms\" by Qijing Huang,\n",
16-
"Po-An Tsai, Joel S. Emer, Angshuman Parashar.\n",
16+
"Po-An Tsai, Joel S. Emer, Angshuman Parashar. The paper includes additional analysis and\n",
17+
"example use cases for this tradeoff.\n",
1718
"\n",
1819
"Our plan to do this analysis is the following:\n",
1920
"- Set up a simple architecture with a main memory and a global buffer\n",
@@ -59,36 +60,27 @@
5960
"metadata": {},
6061
"outputs": [],
6162
"source": [
62-
"# import accelforge as af\n",
63-
"# from pathlib import Path\n",
63+
"import accelforge as af\n",
64+
"from pathlib import Path\n",
6465
"\n",
65-
"# examples_dir = Path(\"../../examples\")\n",
66+
"examples_dir = Path(\"../../examples\")\n",
6667
"\n",
67-
"# spec = af.Spec.from_yaml(\n",
68-
"# af.examples.arches.orojenesis,\n",
69-
"# af.examples.workloads.matmuls,\n",
70-
"# jinja_parse_data={\"N_EINSUMS\": 1, \"M\": 4096, \"KN\": 4096},\n",
71-
"# )\n",
72-
"# spec.mapper.metrics = af.Metrics.ENERGY | af.Metrics.RESOURCE_USAGE"
73-
]
74-
},
75-
{
76-
"cell_type": "code",
77-
"execution_count": null,
78-
"id": "724ba697",
79-
"metadata": {},
80-
"outputs": [],
81-
"source": [
82-
"# # af.set_n_parallel_jobs(1)\n",
83-
"# results = spec.map_workload_to_arch()"
68+
"spec = af.Spec.from_yaml(\n",
69+
" af.examples.arches.orojenesis,\n",
70+
" af.examples.workloads.matmuls,\n",
71+
" jinja_parse_data={\"N_EINSUMS\": 1, \"M\": 4096, \"KN\": 4096},\n",
72+
")\n",
73+
"spec.mapper.metrics = af.Metrics.ENERGY | af.Metrics.RESOURCE_USAGE\n",
74+
"results = spec.map_workload_to_arch()"
8475
]
8576
},
8677
{
8778
"cell_type": "markdown",
8879
"id": "0c06b6c5",
8980
"metadata": {},
9081
"source": [
91-
"Let's plot the results."
82+
"Let's plot the results. We can see that, on a log-log scale, the curve appears to be\n",
83+
"linear."
9284
]
9385
},
9486
{
@@ -98,30 +90,35 @@
9890
"metadata": {},
9991
"outputs": [],
10092
"source": [
101-
"# import matplotlib.pyplot as plt\n",
102-
"\n",
103-
"# results.data.sort_values(\"Total<SEP>energy\", ascending=True, inplace=True)\n",
104-
"# plt.plot(\n",
105-
"# [x * spec.arch.find(\"GlobalBuffer\").size for x in results.resource_usage()[\"GlobalBuffer\"]],\n",
106-
"# results.energy()\n",
107-
"# )\n",
108-
"# plt.xlabel(\"Global Buffer Size (bits)\")\n",
109-
"# plt.ylabel(\"Lowest-Attainable DRAM Accesses (bits)\")\n",
110-
"# plt.xscale(\"log\")\n",
111-
"# plt.yscale(\"log\")\n",
112-
"# plt.show()\n",
113-
"# # Plotting runoff to the right\n",
114-
"# # arxiv and let timeloop team know\n",
115-
"# # restricted imperfect factorization for spatial fanouts"
93+
"import matplotlib.pyplot as plt\n",
94+
"\n",
95+
"results.data.sort_values(\"Total<SEP>energy\", inplace=True)\n",
96+
"fig, ax = plt.subplots(figsize=(10, 10))\n",
97+
"ax.plot(\n",
98+
" [x * spec.arch.find(\"GlobalBuffer\").size for x in results.resource_usage()[\"GlobalBuffer\"]],\n",
99+
" results.energy(),\n",
100+
" marker=\"o\",\n",
101+
" drawstyle=\"steps-pre\"\n",
102+
")\n",
103+
"ax.set_xlabel(\"On-Chip Memory Size (bits)\")\n",
104+
"ax.set_ylabel(\"Lowest-Attainable Off-Chip Accesses (bits)\")\n",
105+
"ax.set_xscale(\"log\")\n",
106+
"ax.set_yscale(\"log\")\n",
107+
"plt.show()"
116108
]
117109
},
118110
{
119111
"cell_type": "markdown",
120112
"id": "5bee35e9",
121113
"metadata": {},
122114
"source": [
115+
"### Fused vs. Unfused\n",
116+
"\n",
123117
"Let's also do a comparison of how the fusion affects this curve. We'll use the same\n",
124-
"architecture, but this time with a larger workload."
118+
"architecture, but this time with a larger workload. We'll run a GPT-3 6.7B attention\n",
119+
"head with 8192 tokens.\n",
120+
"\n",
121+
"Note: Depending on your machine, this may take 3-15 minutes to run."
125122
]
126123
},
127124
{
@@ -143,21 +140,50 @@
143140
" jinja_parse_data={\"N_TOKENS\": 8192}\n",
144141
")\n",
145142
"spec.mapper.metrics = af.Metrics.ENERGY | af.Metrics.RESOURCE_USAGE\n",
146-
"# spec.mapper.max_pmapping_templates_per_einsum = 8\n",
147-
"\n",
148143
"\n",
149144
"# FUSED\n",
150-
"af.set_n_parallel_jobs(1)\n",
145+
"af.set_n_parallel_jobs(32)\n",
151146
"spec.arch.find(\"MainMemory\").tensors.keep = \"~Intermediates\"\n",
152147
"spec.arch.find(\"MainMemory\").tensors.may_keep = \"All\"\n",
153-
"results_fused = spec.map_workload_to_arch(einsum_names=[\"Q\"])\n",
154-
"\n",
148+
"results_fused = spec.map_workload_to_arch()\n",
155149
"\n",
156-
"# # UNFUSED\n",
150+
"# UNFUSED\n",
157151
"spec.arch.find(\"MainMemory\").tensors.keep = \"All\"\n",
158-
"results_unfused = spec.map_workload_to_arch(einsum_names=[\"Q\"])#einsum_names=[\"K\", \"QK\"])\n",
152+
"results_unfused = spec.map_workload_to_arch()"
153+
]
154+
},
155+
{
156+
"cell_type": "markdown",
157+
"id": "dbb84dfd",
158+
"metadata": {},
159+
"source": [
160+
"Let's plot the results. \n",
161+
"\n",
162+
"\n",
163+
"We can see that, for small on-chip memory sizes, there is an approximately-linear\n",
164+
"relationship between on-chip memory size and the number of off-chip accesses.\n",
165+
"Additionally in this region, fusion has negligible benefit. These are because the\n",
166+
"on-chip memory is too small to fully reuse the tensors of any given Einsum, so off-chip\n",
167+
"accesses are dominated by refetching tensors.\n",
159168
"\n",
160-
"# BUG C: Initial stride appearing in the model output when stride == initial"
169+
"As we increase on-chip memory size, the curve begins to flatten out, and fusion begins\n",
170+
"to have a larger effect. These are because the available on-chip memory is large enough\n",
171+
"to store full tensors, which enables multiple fusion opportunities:\n",
172+
"\n",
173+
"- Untiled fusion: We can generate a full output tensor, keep it on-chip, then use it for\n",
174+
" the next Einsum\n",
175+
"- Tiled fusion: We can keep a full weight tensor on-chip, while consuming input tensor\n",
176+
" and sending output tensors one-tile-at-a-time from and to other Einsums without going\n",
177+
" off-chip.\n",
178+
"\n",
179+
"Also note that fusion tends to be less helpful if we're refetching tensors multiple\n",
180+
"times from off-chip, because fusion can save, at most, one fetch of a tensor per Einsum\n",
181+
"(exchanging with another Einsum happens once). If we're refetching tensors multiple\n",
182+
"times, we're likely better off using all available capacity to minimize refetches.\n",
183+
"\n",
184+
"The curve slowly flattens out, rather than quickly, because each Einsum has a different\n",
185+
"shape, and some Einsums enter the no-refetch, fusion-helpful region at smaller on-chip\n",
186+
"memory sizes."
161187
]
162188
},
163189
{
@@ -170,33 +196,39 @@
170196
"results_unfused.resource_usage()\n",
171197
"results_unfused.columns\n",
172198
"import matplotlib.pyplot as plt\n",
199+
"import numpy as np\n",
200+
"\n",
201+
"fig, ax = plt.subplots(figsize=(10, 10))\n",
173202
"\n",
174-
"results_fused.data.sort_values(\"Total<SEP>energy\", ascending=True, inplace=True)\n",
175-
"plt.plot(\n",
176-
" results_fused.resource_usage()[\"GlobalBuffer\"],\n",
203+
"results_fused.data.sort_values(\"Total<SEP>energy\", inplace=True)\n",
204+
"ax.plot(\n",
205+
" np.array(results_fused.resource_usage()[\"GlobalBuffer\"]) * spec.arch.find(\"GlobalBuffer\").size,\n",
177206
" results_fused.energy(),\n",
178207
" label=\"Fused\",\n",
208+
" marker=\"o\",\n",
209+
" drawstyle=\"steps-pre\"\n",
179210
")\n",
180211
"\n",
181-
"results_unfused.data.sort_values(\"Total<SEP>energy\", ascending=True, inplace=True)\n",
182-
"plt.plot(\n",
183-
" results_unfused.resource_usage()[\"GlobalBuffer\"],\n",
212+
"results_unfused.data.sort_values(\"Total<SEP>energy\", inplace=True)\n",
213+
"ax.plot(\n",
214+
" np.array(results_unfused.resource_usage()[\"GlobalBuffer\"]) * spec.arch.find(\"GlobalBuffer\").size,\n",
184215
" results_unfused.energy(),\n",
185216
" label=\"Unfused\",\n",
217+
" marker=\"o\",\n",
218+
" drawstyle=\"steps-pre\"\n",
186219
")\n",
187-
"plt.xlabel(\"Global Buffer Size (bits)\")\n",
188-
"plt.ylabel(\"Lowest-Attainable DRAM Accesses (bits)\")\n",
189-
"plt.xscale(\"log\")\n",
190-
"plt.yscale(\"log\")\n",
191-
"plt.legend()\n",
192-
"plt.show()\n",
193-
"# results_fused.resource_usage(\"GlobalBuffer\")\n"
220+
"ax.set_xlabel(\"On-Chip Memory Size (bits)\")\n",
221+
"ax.set_ylabel(\"Lowest-Attainable Off-Chip Accesses (bits)\")\n",
222+
"ax.set_xscale(\"log\")\n",
223+
"ax.set_yscale(\"log\")\n",
224+
"ax.legend()\n",
225+
"plt.show()"
194226
]
195227
}
196228
],
197229
"metadata": {
198230
"kernelspec": {
199-
"display_name": "venv",
231+
"display_name": "Python 3",
200232
"language": "python",
201233
"name": "python3"
202234
},
@@ -210,7 +242,7 @@
210242
"name": "python",
211243
"nbconvert_exporter": "python",
212244
"pygments_lexer": "ipython3",
213-
"version": "3.12.3"
245+
"version": "3.12.11"
214246
}
215247
},
216248
"nbformat": 4,

0 commit comments

Comments
 (0)