[examples] Add GPT-3 175B workload

Author: gilbertmike

examples/workloads/gpt3_175B.yaml

@@ -0,0 +1,60 @@
+# Each tensor is shaped by a set of ranks, denoted by capital letters
+# For example: Q is shaped by (B, M, H, E)
+# We'll use lower-case letters to index into the ranks
+# For example: Q[b, m, h, e] is the tensor Q at index (b, m, h, e)
+
+# When making a projection list, it's equivalent to the Einsum subscript notation, so:
+# Q projection [b, m, h, e] means that b indexes into B, m indexes into M...
+# When making a projection dict, it's equivalent to the Einsum subscript/superscript notation, so:
+# K projection { B: b, M: p, H: h, E: e } means that b indexes into B, p indexes into M...
+
+# Renames take a tensor name and turn them into a canonical name that we can use in
+# architecture constraints. For example, we want to use the words "input", "weight", and
+# "output" to refer to the tensors of an Einsum, but the Einsum QK has no clear "weight"
+# or "input" because both Q and K are inputs. So we rename K to be weight.
+
+
+workload:
+  rank_sizes:
+    {% set BATCH_SIZE = BATCH_SIZE | default(1) %}
+    {% set N_TOKENS = N_TOKENS | default(8192) %}
+    B: {{BATCH_SIZE}}
+    P: {{N_TOKENS}}
+    M: {{N_TOKENS}}
+    H: 96
+    E: 128
+    F: 128
+    D: 96*128 # = e * h
+    C: 4*96*128
+    J: 96*128
+    G: 96*128
+
+  bits_per_value: {All: 8}
+  persistent_tensors: weight - Intermediates
+
+  einsums:
+  - einsum: I(output)[b, m, d] = I_in(input)[b, m, d]
+    is_copy_operation: True
+  - V[b, m, h, e] = I[b, m, d] * WV[h, e, d]
+  - K[b, m, h, e] = I[b, m, d] * WK[h, e, d]
+  - Q[b, m, h, e] = I[b, m, d] * WQ[h, e, d]
+  - QK[b, m, p, h] = Q(input)[b, m, h, e] * K(weight)[b, M=p, h, e]
+  - QK_softmax[b, m, p, h] = QK(input)[b, m, p, h]
+  - AV[b, m, h, f] = QK_softmax(input)[b, m, p, h] * V(weight)[b, M=p, H=h, E=f]
+  - Z[b, m, g] = AV[b, m, h, f] * WZ[h, f, g]
+  - FFA[b, m, c] = Z[b, m, g] * WFFA[g, c]
+  - FFB[b, m, j] = FFA[b, m, c] * WFFB[c, j]
+
+renames:
+  einsums:
+  - name: default
+    tensor_accesses:
+    - name: input
+      source: Inputs & Intermediates
+      expected_count: 1
+    - name: output
+      source: Outputs
+      expected_count: 1
+    - name: weight
+      source: ~(input | output)
+      expected_count: 1 if len(All) == 3 else 0