Skip to content

Commit 1f1bdbf

Browse files
tanner-andrulistanner-andrulis
committed
Bug notebook
1 parent e6e2ddc commit 1f1bdbf

2 files changed

Lines changed: 250 additions & 0 deletions

File tree

‎examples/arches/oroejenesis.yaml‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# This architecture is used to generate Orojenesis ski-slope that show the tradeoff
2+
# between the size of a memory and the minimum number of accesses to parent memories.
3+
# These trade-offs are introduced in "Mind the Gap: Attainable Data Movement and
4+
# Operational Intensity Bounds for Tensor Algorithms" by Qijing Huang, Po-An Tsai, Joel
5+
# S. Emer, Angshuman Parashar.
6+
7+
# See the memory_size_vs_access_tradeoff.ipynb notebook for more details.
8+
9+
arch:
10+
nodes:
11+
# Set main memory read and write energy to 1J/bit. Then, when we set the mapper to
12+
# minimize energy, it will really just be minimizing the number of DRAM accesses.
13+
- !Memory
14+
name: MainMemory
15+
component_class: Dummy
16+
size: inf
17+
actions: [{name: read, energy: 1}, {name: write, energy: 1}]
18+
# Keep only intermediates -> allow fusion
19+
tensors: {keep: ~Intermediates, may_keep: All}
20+
21+
# We'll set up the mapper to also optimize for global buffer usage. Note: Global
22+
# buffer size can not be infinity (because usage would be forced to zero), but we make
23+
# it large enough that it'll never reasonably be oversubscribed.
24+
- !Memory
25+
name: GlobalBuffer
26+
component_class: Dummy
27+
size: 1e12
28+
tensors: {keep: All}
29+
30+
- !Compute
31+
name: MAC
32+
component_class: Dummy

‎notebooks/tutorials/memory_size_vs_access_tradeoff.ipynb‎

Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,218 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "0fe9944d",
6+
"metadata": {},
7+
"source": [
8+
"## Tradeoffs in Memory Size vs. Accesses to Parent Memories\n",
9+
"\n",
10+
"This notebook will show you how to analyze tradeoffs in memory size versus accesses to\n",
11+
"parent memories. The generated curve will show you, for a given memory size, the lower\n",
12+
"bound of the number of accesses to parent memories.\n",
13+
"\n",
14+
"This analysis is called Orojenesis, and it is introduced in \"Mind the Gap: Attainable\n",
15+
"Data Movement and Operational Intensity Bounds for Tensor Algorithms\" by Qijing Huang,\n",
16+
"Po-An Tsai, Joel S. Emer, Angshuman Parashar.\n",
17+
"\n",
18+
"Our plan to do this analysis is the following:\n",
19+
"- Set up a simple architecture with a main memory and a global buffer\n",
20+
"- Tell the mapper to optimize for both main memory accesses and global buffer usage\n",
21+
"- Observe the resulting Pareto frontier between global buffer usage (size) and the\n",
22+
" minimum number of accesses to main memory.\n",
23+
"\n",
24+
"To this end, we have an \"orojenesis\" architecture in the examples/arches directory.\n",
25+
"Let's take a look at it:\n"
26+
]
27+
},
28+
{
29+
"cell_type": "code",
30+
"execution_count": null,
31+
"id": "fc2a0ee9",
32+
"metadata": {},
33+
"outputs": [],
34+
"source": [
35+
"from IPython.display import Markdown, display\n",
36+
"import accelforge as af\n",
37+
"\n",
38+
"\n",
39+
"display(Markdown(f\"\"\"\n",
40+
"``` yaml\n",
41+
"{open(af.examples.arches.oroejenesis).read()}\n",
42+
"```\n",
43+
"\"\"\"))\n"
44+
]
45+
},
46+
{
47+
"cell_type": "markdown",
48+
"id": "63668ed9",
49+
"metadata": {},
50+
"source": [
51+
"First import AccelForge and initialize a Spec. We'll use a 4096x4096x4096 matrix\n",
52+
"multiply as our workload."
53+
]
54+
},
55+
{
56+
"cell_type": "code",
57+
"execution_count": null,
58+
"id": "5f9af75b",
59+
"metadata": {},
60+
"outputs": [],
61+
"source": [
62+
"# import accelforge as af\n",
63+
"# from pathlib import Path\n",
64+
"\n",
65+
"# examples_dir = Path(\"../../examples\")\n",
66+
"\n",
67+
"# spec = af.Spec.from_yaml(\n",
68+
"# af.examples.arches.oroejenesis,\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()"
84+
]
85+
},
86+
{
87+
"cell_type": "markdown",
88+
"id": "0c06b6c5",
89+
"metadata": {},
90+
"source": [
91+
"Let's plot the results."
92+
]
93+
},
94+
{
95+
"cell_type": "code",
96+
"execution_count": null,
97+
"id": "eff9b08c",
98+
"metadata": {},
99+
"outputs": [],
100+
"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"
116+
]
117+
},
118+
{
119+
"cell_type": "markdown",
120+
"id": "5bee35e9",
121+
"metadata": {},
122+
"source": [
123+
"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."
125+
]
126+
},
127+
{
128+
"cell_type": "code",
129+
"execution_count": null,
130+
"id": "9039d55c",
131+
"metadata": {},
132+
"outputs": [],
133+
"source": [
134+
"import accelforge as af\n",
135+
"from pathlib import Path\n",
136+
"\n",
137+
"examples_dir = Path(\"../../examples\")\n",
138+
"\n",
139+
"spec = af.Spec.from_yaml(\n",
140+
" af.examples.arches.oroejenesis,\n",
141+
" af.examples.workloads.gpt3_6_7B,\n",
142+
" # af.examples.workloads.three_matmuls_annotated,\n",
143+
" jinja_parse_data={\"N_TOKENS\": 8192}\n",
144+
")\n",
145+
"spec.mapper.metrics = af.Metrics.ENERGY | af.Metrics.RESOURCE_USAGE\n",
146+
"# spec.mapper.max_pmapping_templates_per_einsum = 8\n",
147+
"\n",
148+
"\n",
149+
"# FUSED\n",
150+
"af.set_n_parallel_jobs(1)\n",
151+
"spec.arch.find(\"MainMemory\").tensors.keep = \"~Intermediates\"\n",
152+
"spec.arch.find(\"MainMemory\").tensors.may_keep = \"All\"\n",
153+
"results_fused = spec.map_workload_to_arch(einsum_names=[\"Q\"])\n",
154+
"\n",
155+
"\n",
156+
"# # UNFUSED\n",
157+
"spec.arch.find(\"MainMemory\").tensors.keep = \"All\"\n",
158+
"results_unfused = spec.map_workload_to_arch(einsum_names=[\"Q\"])#einsum_names=[\"K\", \"QK\"])\n",
159+
"\n",
160+
"# BUG C: Initial stride appearing in the model output when stride == initial"
161+
]
162+
},
163+
{
164+
"cell_type": "code",
165+
"execution_count": null,
166+
"id": "d9aba6ac",
167+
"metadata": {},
168+
"outputs": [],
169+
"source": [
170+
"results_unfused.resource_usage()\n",
171+
"results_unfused.columns\n",
172+
"import matplotlib.pyplot as plt\n",
173+
"\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",
177+
" results_fused.energy(),\n",
178+
" label=\"Fused\",\n",
179+
")\n",
180+
"\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",
184+
" results_unfused.energy(),\n",
185+
" label=\"Unfused\",\n",
186+
")\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"
194+
]
195+
}
196+
],
197+
"metadata": {
198+
"kernelspec": {
199+
"display_name": "Python 3",
200+
"language": "python",
201+
"name": "python3"
202+
},
203+
"language_info": {
204+
"codemirror_mode": {
205+
"name": "ipython",
206+
"version": 3
207+
},
208+
"file_extension": ".py",
209+
"mimetype": "text/x-python",
210+
"name": "python",
211+
"nbconvert_exporter": "python",
212+
"pygments_lexer": "ipython3",
213+
"version": "3.12.11"
214+
}
215+
},
216+
"nbformat": 4,
217+
"nbformat_minor": 5
218+
}

0 commit comments

Comments
 (0)