MCPcopy Create free account
hub / github.com/apache/tvm / test_direct_sum

Function test_direct_sum

tests/python/tirx/test_layout.py:1294–1328  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

1292
1293
1294def test_direct_sum():
1295 def case1():
1296 # Example from the appendix: A + B yields contiguous (16):(1)
1297 # B = (2,2):(4,1), A = (2,2):(8,2)
1298 B = TileLayout(S[(2, 2) : (4, 1)])
1299 A = TileLayout(S[(2, 2) : (8, 2)])
1300
1301 # Compute direct sum on tiling domain S_A ⊗ S_B with shapes (2,2) and (2,2)
1302 sum_layout = B.direct_sum(A, [2, 2], [2, 2]).canonicalize()
1303 expected = TileLayout(S[16:1])
1304 assert_structural_equal(expected, sum_layout)
1305
1306 # Verify Apply equality: 8p + 2q + 4i + j
1307 print(f"sum_layout: {sum_layout}")
1308 an = Analyzer()
1309 for p in [0, 1]:
1310 for q in [0, 1]:
1311 for i in [0, 1]:
1312 for j in [0, 1]:
1313 m = sum_layout.apply(p, q, i, j, shape=(2, 2, 2, 2))["m"]
1314 m_left = A.apply(p, i, shape=(2, 2))["m"]
1315 m_right = B.apply(q, j, shape=(2, 2))["m"]
1316 assert an.can_prove(m == m_left + m_right)
1317
1318 # Recognition: recover A given B and sum, and recover B given A and sum
1319 interleaved_shape = [2, 2, 2, 2] # [A0, B0, A1, B1]
1320 A_rec = B.is_direct_sum_right(sum_layout, interleaved_shape, [2, 2])
1321 assert A_rec is not None
1322 assert_structural_equal(A.canonicalize(), A_rec.canonicalize())
1323
1324 B_rec = A.is_direct_sum_left(sum_layout, interleaved_shape, [2, 2])
1325 assert B_rec is not None
1326 assert_structural_equal(B.canonicalize(), B_rec.canonicalize())
1327
1328 case1()
1329
1330
1331def test_group_by_logical_shape():

Callers

nothing calls this directly

Calls 1

case1Function · 0.85

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…