[codegen] added fallback when tensor cores cannot be used
This commit is contained in:
@@ -42,7 +42,12 @@ void matmul(restrict read_only align(16) fp16 *A,
|
||||
fp16* pb[TN, TK] = B + rkb[newaxis, :]*ldb + ryb[:, newaxis];
|
||||
fp16 a[TM, TK] = *pa;
|
||||
fp16 b[TN, TK] = *pb;
|
||||
for(int32 k = K; k > TK; k = k - TK){
|
||||
int32 last_a = ((M*K - 1) - (TM*TK + 1)) / lda;
|
||||
int32 last_b = ((K*N - 1) - (TN*TK + 1)) / ldb;
|
||||
last_a = last_a / TK * TK;
|
||||
last_b = last_b / TK * TK;
|
||||
int32 bound = K - max(last_a, last_b);
|
||||
for(int32 k = K; k > bound; k = k - TK){
|
||||
pa = pa + TK*lda;
|
||||
pb = pb + TK*ldb;
|
||||
c = dot(a, trans(b), c);
|
||||
@@ -51,6 +56,15 @@ void matmul(restrict read_only align(16) fp16 *A,
|
||||
}
|
||||
int32 rxc[TM] = get_global_range[TM](0);
|
||||
int32 ryc[TN] = get_global_range[TN](1);
|
||||
for(int32 k = bound; k > 0; k = k - 1){
|
||||
int1 checka[TM, 1] = rxc[:, newaxis] < M;
|
||||
int1 checkb[TN, 1] = ryc[:, newaxis] < N;
|
||||
fp16* pa[TM, 1] = A + (K - k)*lda + rxc[:, newaxis];
|
||||
fp16* pb[TN, 1] = B + (K - k)*ldb + ryc[:, newaxis];
|
||||
fp16 a[TM, 1] = checka ? *pa : 0;
|
||||
fp16 b[TN, 1] = checkb ? *pb : 0;
|
||||
c = dot(a, trans(b), c);
|
||||
}
|
||||
fp32* pc[TM, TN] = C + ryc[newaxis, :]*ldc + rxc[:, newaxis];
|
||||
*pc = c;
|
||||
}
|
||||
|
@@ -6,7 +6,7 @@ data_files_path = tf.resource_loader.get_data_files_path()
|
||||
library_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
module = tf.load_op_library(os.path.join(library_dir, 'libtf_blocksparse.so'))
|
||||
|
||||
M, N, K = 8192, 8192, 8192
|
||||
M, N, K = 128,128,128
|
||||
a = tf.placeholder(tf.float16, shape=[M, K])
|
||||
b = tf.placeholder(tf.float16, shape=[N, K])
|
||||
locks = tf.placeholder(tf.int32, shape=[4096])
|
||||
@@ -24,16 +24,6 @@ result = sess.run([c], feed_dict = {locks: np.zeros(4096),
|
||||
a: ha,
|
||||
b: hb})[0]
|
||||
|
||||
#bench = tf.test.Benchmark().run_op_benchmark(sess=sess,
|
||||
# op_or_tensor=c,
|
||||
# feed_dict={a: ha, b: hb},
|
||||
# min_iters=100)
|
||||
#print(end - start)
|
||||
#print(2*M*N*K / (end - start) * 1e-12)
|
||||
#hresult = np.dot(ha.T, hb).T
|
||||
#dif = np.abs(result - hresult)
|
||||
#print("dif: %f" % np.max(dif))
|
||||
|
||||
#np.savetxt("dif.txt", dif, fmt="%5.2f")
|
||||
#np.savetxt("gpu.txt", result, fmt="%5.2f")
|
||||
#np.savetxt("cpu.txt", hresult, fmt="%5.2f")
|
||||
hresult = np.dot(ha.T, hb).T
|
||||
dif = np.abs(result - hresult)
|
||||
print("dif: %f" % np.max(dif))
|
||||
|
Reference in New Issue
Block a user