add custom MTTKRP to MTTKRP benchmark, turn off mallinfo usage by def…

…ault due to overhead in performance

bench/python/bench_smttkrp.py

@@ -5,7 +5,7 @@
 import sbench_args as sargs
 import numpy as np
 
-def run_bench(num_iter, s_start, s_end, mult, R, sp, sp_init):
+def run_bench(num_iter, s_start, s_end, mult, R, sp, sp_init, use_cust_MTTKRP):
     wrld = ctf.comm()
     s = s_start
     nnz = float(s_start*s_start*s_start)*sp_init
@@ -17,23 +17,28 @@ def run_bench(num_iter, s_start, s_end, mult, R, sp, sp_init):
     agg_max_95 = []
     if num_iter > 1:
         if ctf.comm().rank() == 0:
-            print("Performing MTTKRP WARMUP with s =",s,"nnz =",nnz,"sp =",sp,"sp_init =",sp_init)
+            print("Performing MTTKRP WARMUP with s =",s,"nnz =",nnz,"sp =",sp,"sp_init =",sp_init,"use_cust_MTTKRP=",use_cust_MTTKRP)
         T = ctf.tensor((s,s,s),sp=sp)
         T.fill_sp_random(-1.,1.,float(nnz)/float(s*s*s))
         U = ctf.random.random((s,R))
         V = ctf.random.random((s,R))
         W = ctf.random.random((s,R))
-        U = ctf.einsum("ijk,jr,kr->ir",T,V,W)
-        V = ctf.einsum("ijk,ir,kr->jr",T,U,W)
-        W = ctf.einsum("ijk,ir,jr->kr",T,U,V)
+        if use_cust_MTTKRP:
+            ctf.MTTKRP(T,[U,V,W],0)
+            ctf.MTTKRP(T,[U,V,W],1)
+            ctf.MTTKRP(T,[U,V,W],2)
+        else:
+            U = ctf.einsum("ijk,jr,kr->ir",T,V,W)
+            V = ctf.einsum("ijk,ir,kr->jr",T,U,W)
+            W = ctf.einsum("ijk,ir,jr->kr",T,U,V)
         if ctf.comm().rank() == 0:
-            print("Completed MTTKRP WARMUP with s =",s,"nnz =",nnz,"sp =",sp,"sp_init =",sp_init)
+            print("Completed MTTKRP WARMUP with s =",s,"nnz =",nnz,"sp =",sp,"sp_init =",sp_init,"use_cust_MTTKRP=",use_cust_MTTKRP)
     while s<=s_end:
         agg_s.append(s)
         T = ctf.tensor((s,s,s),sp=sp)
         T.fill_sp_random(-1.,1.,float(nnz)/float(s*s*s))
         if ctf.comm().rank() == 0:
-            print("Performing MTTKRP with s =",s,"nnz =",nnz,"sp =",sp,"sp_init =",sp_init)
+            print("Performing MTTKRP with s =",s,"nnz =",nnz,"sp =",sp,"sp_init =",sp_init,"use_cust_MTTKRP=",use_cust_MTTKRP)
         U = ctf.random.random((s,R))
         V = ctf.random.random((s,R))
         W = ctf.random.random((s,R))
@@ -43,19 +48,28 @@ def run_bench(num_iter, s_start, s_end, mult, R, sp, sp_init):
         avg_times = []
         for i in range(num_iter):
             t0 = time.time()
-            U = ctf.einsum("ijk,jr,kr->ir",T,V,W)
+            if use_cust_MTTKRP == 1:
+                ctf.MTTKRP(T,[U,V,W],0)
+            else:
+                U = ctf.einsum("ijk,jr,kr->ir",T,V,W)
             t1 = time.time()
             ite1 = t1 - t0
             te1 += ite1
 
             t0 = time.time()
-            V = ctf.einsum("ijk,ir,kr->jr",T,U,W)
+            if use_cust_MTTKRP == 1:
+                ctf.MTTKRP(T,[U,V,W],1)
+            else:
+                V = ctf.einsum("ijk,ir,kr->jr",T,U,W)
             t1 = time.time()
             ite2 = t1 - t0
             te2 += ite2
 
             t0 = time.time()
-            W = ctf.einsum("ijk,ir,jr->kr",T,U,V)
+            if use_cust_MTTKRP == 1:
+                ctf.MTTKRP(T,[U,V,W],2)
+            else:
+                W = ctf.einsum("ijk,ir,jr->kr",T,U,V)
             t1 = time.time()
             ite3 = t1 - t0
             te3 += ite3
@@ -66,7 +80,7 @@ def run_bench(num_iter, s_start, s_end, mult, R, sp, sp_init):
             print("Completed",num_iter,"iterations, took",te1/num_iter,te2/num_iter,te3/num_iter,"seconds on average for 3 variants.")
             avg_time = (te1+te2+te3)/(3*num_iter)
             agg_avg_times.append(avg_time)
-            print("MTTKRP took",avg_times,"seconds on average across variants with s =",s,"nnz =",nnz,"sp =",sp,"sp_init =",sp_init)
+            print("MTTKRP took",avg_times,"seconds on average across variants with s =",s,"nnz =",nnz,"sp =",sp,"sp_init =",sp_init,"use_cust_MTTKRP=",use_cust_MTTKRP)
             min_time = np.min(avg_times)
             max_time = np.max(avg_times)
             agg_min_times.append(min_time)
@@ -96,8 +110,9 @@ def run_bench(num_iter, s_start, s_end, mult, R, sp, sp_init):
     R = args.R
     sp = args.sp
     sp_init = args.sp_init
+    use_cust_MTTKRP = args.use_cust_MTTKRP
 
     if ctf.comm().rank() == 0:
-        print("num_iter is",num_iter,"s_start is",s_start,"s_end is",s_end,"mult is",mult,"R is",R,"sp is",sp,"sp_init is",sp_init)
-    run_bench(num_iter, s_start, s_end, mult, R, sp, sp_init)
+        print("num_iter is",num_iter,"s_start is",s_start,"s_end is",s_end,"mult is",mult,"R is",R,"sp is",sp,"sp_init is",sp_init,"use_cust_MTTKRP is",use_cust_MTTKRP)
+    run_bench(num_iter, s_start, s_end, mult, R, sp, sp_init, use_cust_MTTKRP)
 

bench/python/sbench_args.py

@@ -41,3 +41,9 @@ def add_arguments(parser):
         default=1.,
         metavar='float',
         help='Initial sparsity fraction (default: 1)')
+    parser.add_argument(
+        '--use_cust_MTTKRP',
+        type=int,
+        default=1,
+        metavar='int',
+        help='Whether to use MTTKRP routine (default: 1/True)')

src/shared/memcontrol.cxx

@@ -23,11 +23,11 @@
 #include <spi/include/kernel/memory.h>
 #endif
 
-#ifdef __APPLE__
+#ifndef USE_MALLINFO
 #ifndef NOMALLINFO
 #define NOMALLINFO
 #endif
-#endif 
+#endif
 
 #include "../interface/common.h"
 #include "util.h"