()
| 33 | |
| 34 | |
| 35 | def test_prim_func(): |
| 36 | a = tirx.Var("a", "handle") |
| 37 | b = tirx.Var("b", "handle") |
| 38 | func = ( |
| 39 | tirx.PrimFunc( |
| 40 | params=[a, b], |
| 41 | ret_type=None, |
| 42 | buffer_map={ |
| 43 | a: tirx.decl_buffer(shape=[128, 128], dtype="float32", name="A"), |
| 44 | b: tirx.decl_buffer(shape=[256, 256], dtype="float32", name="B"), |
| 45 | }, |
| 46 | body=tirx.Evaluate(0), |
| 47 | ) |
| 48 | .with_attr("global_symbol", "main") |
| 49 | .with_attr("s_tir", True) |
| 50 | ) |
| 51 | _assert_print( |
| 52 | func, |
| 53 | expected=""" |
| 54 | # from tvm.script import tirx as T |
| 55 | # from tvm.tirx.layout import Axis |
| 56 | |
| 57 | @T.prim_func(s_tir=True) |
| 58 | def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")): |
| 59 | T.evaluate(0)""", |
| 60 | ) |
| 61 | |
| 62 | |
| 63 | def test_prim_func_no_sugar_inlined_buffer(): |
nothing calls this directly
no test coverage detected
searching dependent graphs…