bench: arrays.zc, a matmul whose hot loop is proved and a gather that keeps its check

1 files changed, 58 insertions, 0 deletions

diff --git a/tests/bench/arrays.zc b/tests/bench/arrays.zc
new file mode 100644
index 0000000..f83cebe
--- /dev/null
+++ b/tests/bench/arrays.zc
@@ -0,0 +1,58 @@
+namespace;
+include rt;
+
+// two array kernels where the bounds checks are the point, one whose hot loop the
+// compiler proves safe and one that keeps its check on purpose. zy --explain-checks
+// shows which checks stay (suite.py puts the count in the checks column).
+//   m  matmul: c = a * b, 24x24, flat row-major. each row of a and column of b is
+//      copied into a local int[N], and the inner k loop runs k < N over those only,
+//      so its checks are proved. the copies index the flat matrices with i * N + k,
+//      which keeps their checks, but they run N^2 times, not N^3
+//   g  gather: s += data[idx[i]] over 4096 words. idx[i] is proved by i < len(idx);
+//      data[idx[i]] is indexed by a loaded value, so its check stays in the hot loop
+// input: a letter (m, g, or z for neither) and how many times to run it. prints a checksum.
+const int N = 24;
+const int G = 4096;
+
+int matmul(fixed int[?] a, fixed int[?] b, fixed int[?] c) {
+	int[24] ra;
+	int[24] cb;
+	int s = 0;
+	for (int i = 0; i < N; i++) {
+		for (int k = 0; k < N; k++) ra[k] = a[i * N + k];
+		for (int j = 0; j < N; j++) {
+			for (int k = 0; k < N; k++) cb[k] = b[k * N + j];
+			int t = 0;
+			for (int k = 0; k < N; k++) t += ra[k] * cb[k];
+			c[i * N + j] = t;
+			s += t;
+		}
+	}
+	return s;
+}
+int gather(fixed int[?] data, fixed int[?] idx) {
+	int s = 0, n = len(idx);
+	for (int i = 0; i < n; i++) s += data[idx[i]];
+	return s;
+}
+fn main() {
+	let which = rt.get_c(), reps = 0, ch = rt.get_c(), r = 0;
+	while (ch >= '0' && ch <= '9') { reps = reps * 10 + ch - '0'; ch = rt.get_c(); }
+	fixed int[?] a = rt.new(fixed int[?], N * N);
+	fixed int[?] b = rt.new(fixed int[?], N * N);
+	fixed int[?] c = rt.new(fixed int[?], N * N);
+	for (int i = 0; i < N * N; i++) { a[i] = i % 7 - 3; b[i] = i % 5 - 2; }
+	fixed int[?] data = rt.new(fixed int[?], G);
+	fixed int[?] idx = rt.new(fixed int[?], G);
+	// 7919 is odd, so i * 7919 mod 4096 visits every slot once, in a scattered order
+	for (int i = 0; i < G; i++) { data[i] = i; idx[i] = i * 7919 % G; }
+	while (reps > 0) {
+		if (which == 'm') r = matmul(a, b, c);
+		if (which == 'g') r = gather(data, idx);
+		reps--;
+	}
+	rt.put_n(r); rt.put_c(10);
+	rt.free(d*a); rt.free(d*b); rt.free(d*c);
+	rt.free(d*data); rt.free(d*idx);
+	return 0;
+}