diff --git a/bench/python/bench_smttkrp.py b/bench/python/bench_smttkrp.py index 0746ca40..b1ef127a 100644 --- a/bench/python/bench_smttkrp.py +++ b/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) diff --git a/bench/python/sbench_args.py b/bench/python/sbench_args.py index 48e94326..a24aff92 100644 --- a/bench/python/sbench_args.py +++ b/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)') diff --git a/src/shared/memcontrol.cxx b/src/shared/memcontrol.cxx index 6edb2c1e..4289dc2a 100644 --- a/src/shared/memcontrol.cxx +++ b/src/shared/memcontrol.cxx @@ -23,11 +23,11 @@ #include #endif -#ifdef __APPLE__ +#ifndef USE_MALLINFO #ifndef NOMALLINFO #define NOMALLINFO #endif -#endif +#endif #include "../interface/common.h" #include "util.h"