tensor: Fix correctness script

tests/kernel/tensor/check_correctness.py

@@ -1,16 +1,11 @@
 import numpy as np
 import struct
 
-A_array = np.zeros((16, 8))
-B_array = np.zeros((8, 16))
-C_array = np.zeros((16, 16))
-
-file = input("simulator output filename: ")
-
 def hex2float(float_hex_str):
     # print(float_hex_str.strip())
     return struct.unpack(">f",struct.pack(">i",int(float_hex_str,16)))[0]
 
+
 def C_index(threadgroup, thread, register):
     """
     col = ((tg % 4) / 2) * 8;
@@ -43,6 +38,13 @@ def C_index(threadgroup, thread, register):
     return (row, col)
 
 
+def check_sim_output():
+    file = input("simulator output filename: ")
+
+    A_array = np.zeros((16, 8))
+    B_array = np.zeros((8, 16))
+    C_array = np.zeros((16, 16))
+
     with open(file) as f:
         for line in f.readlines():
             line = line.strip()
@@ -80,7 +82,10 @@ with open(file) as f:
                     # print(value)
                     C_array[C_index(threadgroup, thread, register)] = hex2float(value)
 
+    return [A_array, B_array, C_array]
 
+
+if __name__ == "__main__":
     expected = np.load("abc.npz")
     # expected_A = expected['A_array']
     # expected_B = expected['B_array']
@@ -95,13 +100,13 @@ print('expected B:')
     print(expected_B)
     print('expected C:')
     print(expected_C[0:8, 0:8])
+    expected_C.astype('float32').tofile("c_expected.bin")
+
+    [got_A, got_B, got_C] = check_sim_output()
     print('got C:')
     print(C_array[0:8, 0:8])
     print('diff C:')
     print(expected_C[0:8, 0:8] - C_array[0:8, 0:8])
-
+    assert np.allclose(expected_A, got_A)
-expected_C.astype('float32').tofile("c_expected.bin")
+    assert np.allclose(expected_B, got_B)
-
+    assert np.allclose(expected_C, got_C)
-assert np.allclose(expected_A, A_array)
-assert np.allclose(expected_B, B_array)
-assert np.allclose(expected_C, C_array)