matmul32

juno_samples/matmul/src/matmul.jn

@@ -12,3 +12,8 @@ fn matmul<n : usize, m : usize, l : usize>(a : f32[n, m], b : f32[m, l]) -> f32[
 
   return res;
 }
+
+#[entry]
+fn matmul32(a: f32[32, 32], b: f32[32, 32]) -> f32[32, 32] {
+   return matmul::<32, 32, 32>(a, b);
+}
\ No newline at end of file

juno_samples/matmul/src/matmul.sch

@@ -11,16 +11,32 @@ macro simpl!(X) {
   infer-schedules(X);
 }
 
-simpl!(*);
-forkify(*);
-fork-guard-elim(*);
-simpl!(*);
-reduce-slf(*);
-simpl!(*);
-fork-coalesce(*);
-simpl!(*);
-reduce-slf(*);
-simpl!(*);
+if feature("cuda") {
+  inline(matmul32);
+  simpl!(matmul32);
+  let func = outline(matmul32@outer);
+  no-memset(matmul32@res);
+  gpu(func);
+  simpl!(matmul32, func);
+  forkify(func);
+  fork-guard-elim(func);
+  simpl!(func);
+  fork-coalesce(func);
+  fork-coalesce(func);
+  simpl!(func);
+  xdot[true](func);
+}
+
+simpl!(matmul);
+forkify(matmul);
+fork-guard-elim(matmul);
+simpl!(matmul);
+reduce-slf(matmul);
+simpl!(matmul);
+fork-coalesce(matmul);
+simpl!(matmul);
+reduce-slf(matmul);
+simpl!(matmul);
 
 if feature("cuda") {
   // Base