def test_alloc_vtcm(): if not check_prereq_and_setup(): return target = tvm.target.hexagon("v66") buf_len = 2048 A = tvm.te.placeholder((buf_len, ), name="A", dtype="int8") B = tvm.te.placeholder((buf_len, ), name="B", dtype="int8") A_buf = tvm.te.compute((buf_len, ), lambda *i: A(*i), "A_buf") B_buf = tvm.te.compute((buf_len, ), lambda *i: B(*i), "B_buf") C = tvm.te.compute((buf_len, ), lambda *i: A_buf(*i) + B_buf(*i), name="C") s = tvm.te.create_schedule(C.op) # Use VTCM for each buffer. s[A_buf].set_scope("local.vtcm") s[B_buf].set_scope("local.vtcm") config = {"tir.add_lower_pass": hexagon.ir_lower_vtcm_pass()} with tvm.transform.PassContext(config=config): irmod = tvm.lower(s, [A, B, C], name="alloc_vtcm") calls = re.findall("HexagonBackend[A-Za-z]*VTCM", str(irmod["alloc_vtcm"])) assert "HexagonBackendAllocateVTCM" in calls assert "HexagonBackendFreeVTCM" in calls
def test_alloc_vtcm(): if not check_prereq_and_setup(): return target = tvm.target.hexagon('v66') buf_len = 2048 A = tvm.te.placeholder((buf_len, ), name='A', dtype='int8') B = tvm.te.placeholder((buf_len, ), name='B', dtype='int8') A_buf = tvm.te.compute((buf_len, ), lambda *i: A(*i), 'A_buf') B_buf = tvm.te.compute((buf_len, ), lambda *i: B(*i), 'B_buf') C = tvm.te.compute((buf_len, ), lambda *i: A_buf(*i) + B_buf(*i), name='C') s = tvm.te.create_schedule(C.op) # Use VTCM for each buffer. s[A_buf].set_scope("local.vtcm") s[B_buf].set_scope("local.vtcm") config = {'tir.add_lower_pass': hexagon.ir_lower_vtcm_pass()} with tvm.transform.PassContext(config=config): irmod = tvm.lower(s, [A, B, C], name='alloc_vtcm') calls = re.findall('HexagonBackend[A-Za-z]*VTCM', str(irmod['alloc_vtcm'])) assert 'HexagonBackendAllocateVTCM' in calls assert 'HexagonBackendFreeVTCM' in calls