Update sam-kernels to include unfused mttkrp

weiya711 · Jan 22, 2024 · ce48880 · ce48880
1 parent 0b3f6c8
commit ce48880
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/compiler/sam-kernels.sh b/compiler/sam-kernels.sh
@@ -37,6 +37,8 @@ GEN_KERNEL_NAMES=(
   mat_residual
   mat_elemadd3
   tensor3_mttkrp
+  tensor3_mttkrp_unfused1
+  tensor3_mttkrp_unfused2
   vec_spacc_simple
   mat_spacc_simple
   vec_sd_compression_WRONG
@@ -79,6 +81,8 @@ TACO_ARGS=(
   "x(i)=b(i)-C(i,j)*d(j) -f=x:s -f=C:ss -f=b:s -f=d:s"
   "X(i,j)=B(i,j)+C(i,j)+D(i,j) -f=X:ss -f=B:ss -f=C:ss -f=D:ss"
   "X(i,j)=B(i,k,l)*C(j,k)*D(j,l) -f=X:ss -f=B:sss -f=C:ss -f=D:ss" 
+  "T(j,k,l)=C(j,k)*D(j,l) -f=T:sss -f=C:ss -f=D:ss" 
+  "X(i,j)=B(i,k,l)*T(j,k,l) -f=X:ss -f=B:sss -f=T:sss" 
   "x(j)=B(i,j) -f=x:s -f=B:ss"
   "X(j,k)=B(i,j,k) -f=X:ss -f=B:sss"
   "x(i)=b(i) -f=b:s -f=x:d"