From c3aa2e82b99878ae1362ed445fa5a5f961681396 Mon Sep 17 00:00:00 2001 From: Nandeeka Nayak Date: Tue, 23 Jul 2024 19:33:30 -0500 Subject: [PATCH] Add tensor name to constructor --- teaal/trans/header.py | 3 +- tests/integration/demo.yaml | 71 ++++------------------------------ tests/integration/dotprod.py | 2 +- tests/integration/example.py | 4 +- tests/integration/example2.py | 2 +- tests/integration/example3.py | 2 +- tests/integration/example4.py | 2 +- tests/integration/example5.py | 2 +- tests/integration/example6.py | 2 +- tests/integration/example7.py | 4 +- tests/integration/gemm.py | 4 +- tests/integration/gemv.py | 4 +- tests/integration/gram.py | 6 +-- tests/integration/mttkrp.py | 2 +- tests/integration/nrm_sq.py | 4 +- tests/integration/outerprod.py | 2 +- tests/integration/sddmm.py | 4 +- tests/integration/spmm.py | 2 +- tests/integration/spmv.py | 2 +- tests/integration/ttm.py | 2 +- tests/integration/ttv.py | 2 +- tests/trans/test_header.py | 12 +++--- tests/trans/test_hifiber.py | 36 ++++++++--------- 23 files changed, 60 insertions(+), 116 deletions(-) diff --git a/teaal/trans/header.py b/teaal/trans/header.py index 70007f0..e7cd078 100644 --- a/teaal/trans/header.py +++ b/teaal/trans/header.py @@ -99,7 +99,8 @@ def make_output(self) -> Statement: self.program.get_loop_order().apply(tensor) arg0 = TransUtils.build_rank_ids(tensor) - args = self.__make_shape([arg0]) + arg1 = AParam("name", EString(tensor.root_name())) + args = self.__make_shape([arg0, arg1]) constr = EFunc("Tensor", args) return SAssign(AVar(tensor.tensor_name()), constr) diff --git a/tests/integration/demo.yaml b/tests/integration/demo.yaml index 72e2543..35ddd45 100644 --- a/tests/integration/demo.yaml +++ b/tests/integration/demo.yaml @@ -1,66 +1,9 @@ einsum: - declaration: - SOB: [UA, UB] - T: [UA, UB, K] - I: [K] - OB: [] + declaration: # Ranks are listed alphabetically in this section + TS: [T, P1, P0, E] + ND: [P1, P0, E] expressions: - - SOB[ua, ub] = T[ua, ub, k] * I[k] - - OB[] = SOB[ua, ub] -# einsum: -# declaration: -# A: [S] -# Z: [T] -# expressions: -# - Z[t] = A[2 * t] -# format: -# A: -# default: -# rank-order: [S] -# S: -# format: C -# pbits: 32 -# Z: -# default: -# rank-order: [T] -# T: -# format: C -# pbits: 32 -# architecture: -# accel: -# - name: System -# local: -# - name: MainMemory -# class: DRAM -# subtree: -# - name: Chip -# local: -# - name: LLB -# class: Cache -# attributes: -# width: 32 -# depth: 1024 -# bindings: -# Z: -# - config: accel -# prefix: tmp/demo -# - component: MainMemory -# bindings: -# - tensor: A -# rank: S -# type: payload -# format: default -# - tensor: Z -# rank: T -# type: payload -# format: default -# - component: LLB -# bindings: -# - tensor: A -# rank: S -# type: payload -# format: default -# - tensor: Z -# rank: T -# type: payload -# format: default + - ND[p1, p0, e] = TS[t, p1, p0, e] +mapping: + loop-order: + ND: [T, P1, P0, E] diff --git a/tests/integration/dotprod.py b/tests/integration/dotprod.py index 00f844b..d1431c8 100644 --- a/tests/integration/dotprod.py +++ b/tests/integration/dotprod.py @@ -1,4 +1,4 @@ -Z_ = Tensor(rank_ids=[]) +Z_ = Tensor(rank_ids=[], name="Z") z_ref = Z_.getRoot() a_k = A_K.getRoot() b_k = B_K.getRoot() diff --git a/tests/integration/example.py b/tests/integration/example.py index 62d812d..ec4fa74 100644 --- a/tests/integration/example.py +++ b/tests/integration/example.py @@ -1,4 +1,4 @@ -T1_IJ = Tensor(rank_ids=["I", "J"]) +T1_IJ = Tensor(rank_ids=["I", "J"], name="T1") t1_i = T1_IJ.getRoot() a_i = A_IJK.getRoot() b_k = B_KL.getRoot() @@ -7,7 +7,7 @@ for k, (a_val, b_l) in a_k & b_k: for l, b_val in b_l: t1_ref += a_val * b_val -D_I = Tensor(rank_ids=["I"]) +D_I = Tensor(rank_ids=["I"], name="D") d_i = D_I.getRoot() c_i = C_IJ.getRoot() t1_i = T1_IJ.getRoot() diff --git a/tests/integration/example2.py b/tests/integration/example2.py index b8ac778..0fc91e2 100644 --- a/tests/integration/example2.py +++ b/tests/integration/example2.py @@ -1,4 +1,4 @@ -D_ = Tensor(rank_ids=[]) +D_ = Tensor(rank_ids=[], name="D") B_IJ = B_JI.swizzleRanks(rank_ids=["I", "J"]) d_ref = D_.getRoot() a_i = A_IJ.getRoot() diff --git a/tests/integration/example3.py b/tests/integration/example3.py index b246d75..80089ef 100644 --- a/tests/integration/example3.py +++ b/tests/integration/example3.py @@ -1,4 +1,4 @@ -D_ = Tensor(rank_ids=[]) +D_ = Tensor(rank_ids=[], name="D") B_IJK = B_JKI.swizzleRanks(rank_ids=["I", "J", "K"]) d_ref = D_.getRoot() a_i = A_IJK.getRoot() diff --git a/tests/integration/example4.py b/tests/integration/example4.py index 330c7f5..cdcbc7e 100644 --- a/tests/integration/example4.py +++ b/tests/integration/example4.py @@ -1,4 +1,4 @@ -D_ = Tensor(rank_ids=[]) +D_ = Tensor(rank_ids=[], name="D") B_IJK = B_JKI.swizzleRanks(rank_ids=["I", "J", "K"]) C_IJK = C_JKI.swizzleRanks(rank_ids=["I", "J", "K"]) d_ref = D_.getRoot() diff --git a/tests/integration/example5.py b/tests/integration/example5.py index b25295d..6b8990d 100644 --- a/tests/integration/example5.py +++ b/tests/integration/example5.py @@ -1,4 +1,4 @@ -Z_IK = Tensor(rank_ids=["I", "K"]) +Z_IK = Tensor(rank_ids=["I", "K"], name="Z") B_KJ = B_JK.swizzleRanks(rank_ids=["K", "J"]) z_i = Z_IK.getRoot() a_i = A_IJ.getRoot() diff --git a/tests/integration/example6.py b/tests/integration/example6.py index 27fdbb9..fed3b07 100644 --- a/tests/integration/example6.py +++ b/tests/integration/example6.py @@ -1,4 +1,4 @@ -D_I = Tensor(rank_ids=["I"]) +D_I = Tensor(rank_ids=["I"], name="D") d_i = D_I.getRoot() c_i = C_IJ.getRoot() for i, (d_ref, c_j) in d_i << c_i: diff --git a/tests/integration/example7.py b/tests/integration/example7.py index 84d413d..1264b61 100644 --- a/tests/integration/example7.py +++ b/tests/integration/example7.py @@ -1,11 +1,11 @@ -Z_MN = Tensor(rank_ids=["M", "N"]) +Z_MN = Tensor(rank_ids=["M", "N"], name="Z") z_m = Z_MN.getRoot() a_m = A_MN.getRoot() b_m = B_MN.getRoot() for m, (z_n, (_, a_n, b_n)) in z_m << (a_m | b_m): for n, (z_ref, (_, a_val, b_val)) in z_n << (a_n | b_n): z_ref += a_val + b_val -Z_MN = Tensor(rank_ids=["M", "N"]) +Z_MN = Tensor(rank_ids=["M", "N"], name="Z") z_m = Z_MN.getRoot() a_m = A_MN.getRoot() b_m = B_MN.getRoot() diff --git a/tests/integration/gemm.py b/tests/integration/gemm.py index b601168..c4d1c43 100644 --- a/tests/integration/gemm.py +++ b/tests/integration/gemm.py @@ -1,4 +1,4 @@ -T1_MN = Tensor(rank_ids=["M", "N"]) +T1_MN = Tensor(rank_ids=["M", "N"], name="T1") A_MK = A_KM.swizzleRanks(rank_ids=["M", "K"]) B_NK = B_KN.swizzleRanks(rank_ids=["N", "K"]) t1_m = T1_MN.getRoot() @@ -8,7 +8,7 @@ for n, (t1_ref, b_k) in t1_n << b_n: for k, (a_val, b_val) in a_k & b_k: t1_ref += a_val * b_val -Z_MN = Tensor(rank_ids=["M", "N"]) +Z_MN = Tensor(rank_ids=["M", "N"], name="Z") z_m = Z_MN.getRoot() t1_m = T1_MN.getRoot() c_m = C_MN.getRoot() diff --git a/tests/integration/gemv.py b/tests/integration/gemv.py index 205dcb6..287dcda 100644 --- a/tests/integration/gemv.py +++ b/tests/integration/gemv.py @@ -1,4 +1,4 @@ -T1_M = Tensor(rank_ids=["M"]) +T1_M = Tensor(rank_ids=["M"], name="T1") B_MK = B_KM.swizzleRanks(rank_ids=["M", "K"]) t1_m = T1_M.getRoot() a_k = A_K.getRoot() @@ -6,7 +6,7 @@ for m, (t1_ref, b_k) in t1_m << b_m: for k, (a_val, b_val) in a_k & b_k: t1_ref += a_val * b_val -Z_M = Tensor(rank_ids=["M"]) +Z_M = Tensor(rank_ids=["M"], name="Z") z_m = Z_M.getRoot() t1_m = T1_M.getRoot() c_m = C_M.getRoot() diff --git a/tests/integration/gram.py b/tests/integration/gram.py index 9f2ad46..d8b5f45 100644 --- a/tests/integration/gram.py +++ b/tests/integration/gram.py @@ -1,4 +1,4 @@ -G0_II_ = Tensor(rank_ids=["I", "I_"]) +G0_II_ = Tensor(rank_ids=["I", "I_"], name="G0") g0_i = G0_II_.getRoot() a_i = A_IJK.getRoot() b0_i_ = B0_I_JK.getRoot() @@ -7,7 +7,7 @@ for j, (a_k, b0_k) in a_j & b0_j: for k, (a_val, b0_val) in a_k & b0_k: g0_ref += a_val * b0_val -G1_JJ_ = Tensor(rank_ids=["J", "J_"]) +G1_JJ_ = Tensor(rank_ids=["J", "J_"], name="G1") A_JIK = A_IJK.swizzleRanks(rank_ids=["J", "I", "K"]) B1_J_IK = B1_IJ_K.swizzleRanks(rank_ids=["J_", "I", "K"]) g1_j = G1_JJ_.getRoot() @@ -18,7 +18,7 @@ for i, (a_k, b1_k) in a_i & b1_i: for k, (a_val, b1_val) in a_k & b1_k: g1_ref += a_val * b1_val -G2_KK_ = Tensor(rank_ids=["K", "K_"]) +G2_KK_ = Tensor(rank_ids=["K", "K_"], name="G2") A_KIJ = A_IJK.swizzleRanks(rank_ids=["K", "I", "J"]) B2_K_IJ = B2_IJK_.swizzleRanks(rank_ids=["K_", "I", "J"]) g2_k = G2_KK_.getRoot() diff --git a/tests/integration/mttkrp.py b/tests/integration/mttkrp.py index ee8663b..5d8391a 100644 --- a/tests/integration/mttkrp.py +++ b/tests/integration/mttkrp.py @@ -1,4 +1,4 @@ -Z_MN = Tensor(rank_ids=["M", "N"]) +Z_MN = Tensor(rank_ids=["M", "N"], name="Z") B_NK = B_KN.swizzleRanks(rank_ids=["N", "K"]) C_NJ = C_JN.swizzleRanks(rank_ids=["N", "J"]) z_m = Z_MN.getRoot() diff --git a/tests/integration/nrm_sq.py b/tests/integration/nrm_sq.py index c4307f4..e0d9ded 100644 --- a/tests/integration/nrm_sq.py +++ b/tests/integration/nrm_sq.py @@ -1,4 +1,4 @@ -T_ABIJ = Tensor(rank_ids=["A", "B", "I", "J"]) +T_ABIJ = Tensor(rank_ids=["A", "B", "I", "J"], name="T") t_a = T_ABIJ.getRoot() v_a = V_ABIJ.getRoot() for a, (t_b, v_b) in t_a << v_a: @@ -6,7 +6,7 @@ for i, (t_j, v_j) in t_i << v_i: for j, (t_ref, v_val) in t_j << v_j: t_ref += v_val -Q_ = Tensor(rank_ids=[]) +Q_ = Tensor(rank_ids=[], name="Q") q_ref = Q_.getRoot() v_a = V_ABIJ.getRoot() t_a = T_ABIJ.getRoot() diff --git a/tests/integration/outerprod.py b/tests/integration/outerprod.py index 72eea82..4c6d6ba 100644 --- a/tests/integration/outerprod.py +++ b/tests/integration/outerprod.py @@ -1,4 +1,4 @@ -Z_MN = Tensor(rank_ids=["M", "N"]) +Z_MN = Tensor(rank_ids=["M", "N"], name="Z") z_m = Z_MN.getRoot() a_m = A_M.getRoot() b_n = B_N.getRoot() diff --git a/tests/integration/sddmm.py b/tests/integration/sddmm.py index 560427d..3a656a0 100644 --- a/tests/integration/sddmm.py +++ b/tests/integration/sddmm.py @@ -1,4 +1,4 @@ -T1_MN = Tensor(rank_ids=["M", "N"]) +T1_MN = Tensor(rank_ids=["M", "N"], name="T1") B_NK = B_KN.swizzleRanks(rank_ids=["N", "K"]) t1_m = T1_MN.getRoot() a_m = A_MK.getRoot() @@ -7,7 +7,7 @@ for n, (t1_ref, b_k) in t1_n << b_n: for k, (a_val, b_val) in a_k & b_k: t1_ref += a_val * b_val -Z_MN = Tensor(rank_ids=["M", "N"]) +Z_MN = Tensor(rank_ids=["M", "N"], name="Z") z_m = Z_MN.getRoot() c_m = C_MN.getRoot() t1_m = T1_MN.getRoot() diff --git a/tests/integration/spmm.py b/tests/integration/spmm.py index e6e3f52..801aac3 100644 --- a/tests/integration/spmm.py +++ b/tests/integration/spmm.py @@ -1,4 +1,4 @@ -Z_MN = Tensor(rank_ids=["M", "N"]) +Z_MN = Tensor(rank_ids=["M", "N"], name="Z") A_MK = A_KM.swizzleRanks(rank_ids=["M", "K"]) B_NK = B_KN.swizzleRanks(rank_ids=["N", "K"]) z_m = Z_MN.getRoot() diff --git a/tests/integration/spmv.py b/tests/integration/spmv.py index af0c69a..78b7f86 100644 --- a/tests/integration/spmv.py +++ b/tests/integration/spmv.py @@ -1,4 +1,4 @@ -Z_M = Tensor(rank_ids=["M"]) +Z_M = Tensor(rank_ids=["M"], name="Z") z_m = Z_M.getRoot() a_k = A_K.getRoot() b_m = B_MK.getRoot() diff --git a/tests/integration/ttm.py b/tests/integration/ttm.py index 5aa4644..9857af0 100644 --- a/tests/integration/ttm.py +++ b/tests/integration/ttm.py @@ -1,4 +1,4 @@ -Z_MNO = Tensor(rank_ids=["M", "N", "O"]) +Z_MNO = Tensor(rank_ids=["M", "N", "O"], name="Z") B_OK = B_KO.swizzleRanks(rank_ids=["O", "K"]) z_m = Z_MNO.getRoot() a_m = A_MNK.getRoot() diff --git a/tests/integration/ttv.py b/tests/integration/ttv.py index e5de231..8b27317 100644 --- a/tests/integration/ttv.py +++ b/tests/integration/ttv.py @@ -1,4 +1,4 @@ -Z_MN = Tensor(rank_ids=["M", "N"]) +Z_MN = Tensor(rank_ids=["M", "N"], name="Z") z_m = Z_MN.getRoot() a_m = A_MNK.getRoot() b_k = B_K.getRoot() diff --git a/tests/trans/test_header.py b/tests/trans/test_header.py index 6c6da39..04cc4b1 100644 --- a/tests/trans/test_header.py +++ b/tests/trans/test_header.py @@ -124,7 +124,7 @@ def test_make_output(): Z: [M1, K, N, M0] """ - hifiber = "Z_M1NM0 = Tensor(rank_ids=[\"M1\", \"N\", \"M0\"])" + hifiber = "Z_M1NM0 = Tensor(rank_ids=[\"M1\", \"N\", \"M0\"], name=\"Z\")" header = build_matmul_header(mapping) assert header.make_output().gen(depth=0) == hifiber @@ -135,7 +135,7 @@ def test_make_output_shape(): - Z[m, n] = A[k, m] """ - hifiber = "Z_MN = Tensor(rank_ids=[\"M\", \"N\"], shape=[M, N])" + hifiber = "Z_MN = Tensor(rank_ids=[\"M\", \"N\"], name=\"Z\", shape=[M, N])" header = build_header(exprs, "") assert header.make_output().gen(depth=0) == hifiber @@ -152,27 +152,27 @@ def test_make_output_no_shape_flattening(): (M, N): [flatten()] """ - hifiber = "Z_MN = Tensor(rank_ids=[\"MN\"])" + hifiber = "Z_MN = Tensor(rank_ids=[\"MN\"], name=\"Z\")" header = build_header(exprs, mapping) assert header.make_output().gen(depth=0) == hifiber def test_make_output_conv_no_shape(): - hifiber = "O_Q = Tensor(rank_ids=[\"Q\"])" + hifiber = "O_Q = Tensor(rank_ids=[\"Q\"], name=\"O\")" header = build_header_conv("[S, Q]") assert header.make_output().gen(0) == hifiber def test_make_output_conv_shape(): - hifiber = "O_Q = Tensor(rank_ids=[\"Q\"], shape=[Q])" + hifiber = "O_Q = Tensor(rank_ids=[\"Q\"], name=\"O\", shape=[Q])" header = build_header_conv("[Q, S]") assert header.make_output().gen(0) == hifiber def test_make_output_metrics_shape(): - hifiber = "T_MKN = Tensor(rank_ids=[\"M\", \"K\", \"N\"], shape=[M, K, N])" + hifiber = "T_MKN = Tensor(rank_ids=[\"M\", \"K\", \"N\"], name=\"T\", shape=[M, K, N])" header = build_header_gamma() assert header.make_output().gen(0) == hifiber diff --git a/tests/trans/test_hifiber.py b/tests/trans/test_hifiber.py index 47b3de3..fa32644 100644 --- a/tests/trans/test_hifiber.py +++ b/tests/trans/test_hifiber.py @@ -6,7 +6,7 @@ def test_translate_no_loops(): einsum = Einsum.from_file("tests/integration/test_translate_no_loops.yaml") mapping = Mapping.from_file( "tests/integration/test_translate_no_loops.yaml") - hifiber = "A_ = Tensor(rank_ids=[])\n" + \ + hifiber = "A_ = Tensor(rank_ids=[], name=\"A\")\n" + \ "a_ref = A_.getRoot()\n" + \ "a_ref += b" assert str(HiFiber(einsum, mapping)) == hifiber @@ -16,7 +16,7 @@ def test_translate_defaults(): einsum = Einsum.from_file("tests/integration/test_input_no_mapping.yaml") mapping = Mapping.from_file( "tests/integration/test_input_no_mapping.yaml") - hifiber = "T1_MN = Tensor(rank_ids=[\"M\", \"N\"])\n" + \ + hifiber = "T1_MN = Tensor(rank_ids=[\"M\", \"N\"], name=\"T1\")\n" + \ "A_MK = A_KM.swizzleRanks(rank_ids=[\"M\", \"K\"])\n" + \ "B_NK = B_KN.swizzleRanks(rank_ids=[\"N\", \"K\"])\n" + \ "t1_m = T1_MN.getRoot()\n" + \ @@ -26,7 +26,7 @@ def test_translate_defaults(): " for n, (t1_ref, b_k) in t1_n << b_n:\n" + \ " for k, (a_val, b_val) in a_k & b_k:\n" + \ " t1_ref += a_val * b_val\n" + \ - "Z_MN = Tensor(rank_ids=[\"M\", \"N\"])\n" + \ + "Z_MN = Tensor(rank_ids=[\"M\", \"N\"], name=\"Z\")\n" + \ "z_m = Z_MN.getRoot()\n" + \ "t1_m = T1_MN.getRoot()\n" + \ "c_m = C_MN.getRoot()\n" + \ @@ -41,7 +41,7 @@ def test_translate_specified(): einsum = Einsum.from_file("tests/integration/test_input.yaml") mapping = Mapping.from_file("tests/integration/test_input.yaml") - hifiber_option1 = "T1_NM = Tensor(rank_ids=[\"N\", \"M\"])\n" + \ + hifiber_option1 = "T1_NM = Tensor(rank_ids=[\"N\", \"M\"], name=\"T1\")\n" + \ "A_KM = A_MK.swizzleRanks(rank_ids=[\"K\", \"M\"])\n" + \ "canvas = createCanvas(A_KM, B_KN, T1_NM)\n" + \ "t1_n = T1_NM.getRoot()\n" + \ @@ -56,7 +56,7 @@ def test_translate_specified(): "tmp1 = tmp0.swizzleRanks(rank_ids=[\"M\", \"N\"])\n" + \ "T1_MN = tmp1\n" + \ "displayCanvas(canvas)\n" + \ - "Z_M2N2M1N1M0N0 = Tensor(rank_ids=[\"M2\", \"N2\", \"M1\", \"N1\", \"M0\", \"N0\"])\n" + \ + "Z_M2N2M1N1M0N0 = Tensor(rank_ids=[\"M2\", \"N2\", \"M1\", \"N1\", \"M0\", \"N0\"], name=\"T1\")\n" + \ "tmp2 = T1_MN\n" + \ "tmp3 = tmp2.splitUniform(6, depth=1)\n" + \ "tmp4 = tmp3.splitUniform(3, depth=2)\n" + \ @@ -96,7 +96,7 @@ def test_translate_specified(): "tmp17.setRankIds(rank_ids=[\"N\", \"M\"])\n" + \ "Z_NM = tmp17" - hifiber_option2 = "T1_NM = Tensor(rank_ids=[\"N\", \"M\"])\n" + \ + hifiber_option2 = "T1_NM = Tensor(rank_ids=[\"N\", \"M\"], name=\"T1\")\n" + \ "A_KM = A_MK.swizzleRanks(rank_ids=[\"K\", \"M\"])\n" + \ "canvas = createCanvas(A_KM, B_KN, T1_NM)\n" + \ "t1_n = T1_NM.getRoot()\n" + \ @@ -111,7 +111,7 @@ def test_translate_specified(): "tmp1 = tmp0.swizzleRanks(rank_ids=[\"M\", \"N\"])\n" + \ "T1_MN = tmp1\n" + \ "displayCanvas(canvas)\n" + \ - "Z_M2N2M1N1M0N0 = Tensor(rank_ids=[\"M2\", \"N2\", \"M1\", \"N1\", \"M0\", \"N0\"])\n" + \ + "Z_M2N2M1N1M0N0 = Tensor(rank_ids=[\"M2\", \"N2\", \"M1\", \"N1\", \"M0\", \"N0\"], name=\"Z\")\n" + \ "tmp2 = T1_MN\n" + \ "tmp3 = tmp2.splitUniform(4, depth=0)\n" + \ "tmp4 = tmp3.splitUniform(2, depth=1)\n" + \ @@ -151,7 +151,7 @@ def test_translate_specified(): "tmp17.setRankIds(rank_ids=[\"N\", \"M\"])\n" + \ "Z_NM = tmp17" - hifiber_option3 = "T1_NM = Tensor(rank_ids=[\"N\", \"M\"])\n" + \ + hifiber_option3 = "T1_NM = Tensor(rank_ids=[\"N\", \"M\"], name=\"T1\")\n" + \ "A_KM = A_MK.swizzleRanks(rank_ids=[\"K\", \"M\"])\n" + \ "t1_n = T1_NM.getRoot()\n" + \ "a_k = A_KM.getRoot()\n" + \ @@ -166,7 +166,7 @@ def test_translate_specified(): "tmp1 = tmp0.swizzleRanks(rank_ids=[\"M\", \"N\"])\n" + \ "T1_MN = tmp1\n" + \ "displayCanvas(canvas)\n" + \ - "Z_M2N2M1N1M0N0 = Tensor(rank_ids=[\"M2\", \"N2\", \"M1\", \"N1\", \"M0\", \"N0\"])\n" + \ + "Z_M2N2M1N1M0N0 = Tensor(rank_ids=[\"M2\", \"N2\", \"M1\", \"N1\", \"M0\", \"N0\"], name=\"Z\")\n" + \ "tmp2 = T1_MN\n" + \ "tmp3 = tmp2.splitUniform(4, depth=0)\n" + \ "tmp4 = tmp3.splitUniform(2, depth=1)\n" + \ @@ -206,7 +206,7 @@ def test_translate_specified(): "tmp17.setRankIds(rank_ids=[\"N\", \"M\"])\n" + \ "Z_NM = tmp17" - hifiber_option4 = "T1_NM = Tensor(rank_ids=[\"N\", \"M\"])\n" + \ + hifiber_option4 = "T1_NM = Tensor(rank_ids=[\"N\", \"M\"], name=\"T1\")\n" + \ "A_KM = A_MK.swizzleRanks(rank_ids=[\"K\", \"M\"])\n" + \ "t1_n = T1_NM.getRoot()\n" + \ "a_k = A_KM.getRoot()\n" + \ @@ -221,7 +221,7 @@ def test_translate_specified(): "tmp1 = tmp0.swizzleRanks(rank_ids=[\"M\", \"N\"])\n" + \ "T1_MN = tmp1\n" + \ "displayCanvas(canvas)\n" + \ - "Z_M2N2M1N1M0N0 = Tensor(rank_ids=[\"M2\", \"N2\", \"M1\", \"N1\", \"M0\", \"N0\"])\n" + \ + "Z_M2N2M1N1M0N0 = Tensor(rank_ids=[\"M2\", \"N2\", \"M1\", \"N1\", \"M0\", \"N0\"], name=\"Z\")\n" + \ "tmp2 = T1_MN\n" + \ "tmp3 = tmp2.splitUniform(6, depth=1)\n" + \ "tmp4 = tmp3.splitUniform(3, depth=2)\n" + \ @@ -290,7 +290,7 @@ def test_hifiber_dyn_part(): einsum = Einsum.from_str(yaml) mapping = Mapping.from_str(yaml) - hifiber = "Z_MN1N0 = Tensor(rank_ids=[\"M\", \"N1\", \"N0\"])\n" + \ + hifiber = "Z_MN1N0 = Tensor(rank_ids=[\"M\", \"N1\", \"N0\"], name=\"Z\")\n" + \ "z_m = Z_MN1N0.getRoot()\n" + \ "a_k = A_KM.getRoot()\n" + \ "b_k = B_KN.getRoot()\n" + \ @@ -361,7 +361,7 @@ def test_hifiber_index_math_no_halo(): einsum = Einsum.from_str(yaml) mapping = Mapping.from_str(yaml) - hifiber = "Z_M2M1M0 = Tensor(rank_ids=[\"M2\", \"M1\", \"M0\"])\n" + \ + hifiber = "Z_M2M1M0 = Tensor(rank_ids=[\"M2\", \"M1\", \"M0\"], name=\"Z\")\n" + \ "tmp0 = A_K\n" + \ "tmp1 = tmp0.splitUniform(2 * 10, depth=0)\n" + \ "tmp2 = tmp1.splitUniform(2 * 5, depth=1)\n" + \ @@ -411,7 +411,7 @@ def test_hifiber_conv(): einsum = Einsum.from_str(yaml) mapping = Mapping.from_str(yaml) - hifiber = "O_Q1Q0 = Tensor(rank_ids=[\"Q1\", \"Q0\"])\n" + \ + hifiber = "O_Q1Q0 = Tensor(rank_ids=[\"Q1\", \"Q0\"], name=\"O\")\n" + \ "tmp0 = I_W\n" + \ "tmp1 = tmp0.splitUniform(10, depth=0, post_halo=-1 + S)\n" + \ "I_W1W0 = tmp1\n" + \ @@ -461,7 +461,7 @@ def test_hifiber_static_flattening(): einsum = Einsum.from_str(yaml) mapping = Mapping.from_str(yaml) - hifiber = "Z_NM = Tensor(rank_ids=[\"N\", \"M\"])\n" + \ + hifiber = "Z_NM = Tensor(rank_ids=[\"N\", \"M\"], name=\"Z\")\n" + \ "tmp0 = A_KM\n" + \ "tmp1 = tmp0.splitUniform(4, depth=0)\n" + \ "A_K1K0M = tmp1\n" + \ @@ -521,7 +521,7 @@ def test_hifiber_dyn_flattening(): einsum = Einsum.from_str(yaml) mapping = Mapping.from_str(yaml) - hifiber = "Z_M1NM0 = Tensor(rank_ids=[\"M1\", \"N\", \"M0\"])\n" + \ + hifiber = "Z_M1NM0 = Tensor(rank_ids=[\"M1\", \"N\", \"M0\"], name=\"Z\")\n" + \ "tmp0 = A_KM\n" + \ "tmp1 = tmp0.splitUniform(6, depth=1)\n" + \ "A_KM1M0 = tmp1\n" + \ @@ -635,7 +635,7 @@ def test_hifiber_traffic(): bindings = Bindings.from_str(yaml) format_ = Format.from_str(yaml) - hifiber = "Z_M = Tensor(rank_ids=[\"M\"], shape=[M])\n" + \ + hifiber = "Z_M = Tensor(rank_ids=[\"M\"], name=\"Z\", shape=[M])\n" + \ "A_MK = A_KM.swizzleRanks(rank_ids=[\"M\", \"K\"])\n" + \ "B_MK = B_KM.swizzleRanks(rank_ids=[\"M\", \"K\"])\n" + \ "z_m = Z_M.getRoot()\n" + \ @@ -711,7 +711,7 @@ def test_hifiber_intersect(): bindings = Bindings.from_str(yaml) format_ = Format.from_str(yaml) - hifiber = "Z_ = Tensor(rank_ids=[], shape=[])\n" + \ + hifiber = "Z_ = Tensor(rank_ids=[], name=\"Z\", shape=[])\n" + \ "z_ref = Z_.getRoot()\n" + \ "a_i = A_IJK.getRoot()\n" + \ "b_i = B_IJK.getRoot()\n" + \