| 333 | |
| 334 | |
| 335 | def test_tuple_get_root(): |
| 336 | def before(dim: int): |
| 337 | bb = relax.BlockBuilder() |
| 338 | x = relax.Var("x", R.Tensor((1, 3 * dim), "float32")) |
| 339 | w = relax.Var("w", R.Tensor((dim, dim), "float32")) |
| 340 | with bb.function("main", [x, w]): |
| 341 | with bb.dataflow(): |
| 342 | lv0 = bb.emit_te(topi.split, x, indices_or_sections=3, axis=1) |
| 343 | lv1 = bb.emit(relax.TupleGetItem(lv0, 0)) |
| 344 | gv = bb.emit_output(bb.call_te(topi.nn.dense, lv1, w)) |
| 345 | bb.emit_func_output(gv) |
| 346 | |
| 347 | return bb.get() |
| 348 | |
| 349 | def expected(dim: int): |
| 350 | bb = relax.BlockBuilder() |
| 351 | |
| 352 | # Grouped function |
| 353 | x = relax.Var("x", R.Tensor((1, 3 * dim), "float32")) |
| 354 | with bb.function("fused_split", [x], attrs={"Primitive": True}, private=True): |
| 355 | with bb.dataflow(): |
| 356 | lv0 = bb.emit_te(topi.split, x, indices_or_sections=3, axis=1) |
| 357 | gv = bb.emit_output(relax.TupleGetItem(lv0, 0)) |
| 358 | bb.emit_func_output(gv) |
| 359 | |
| 360 | # Get the global variables of the grouped functions |
| 361 | fused_split = bb.get().get_global_var("fused_split") |
| 362 | |
| 363 | # Main function |
| 364 | x = relax.Var("x", R.Tensor((1, 3 * dim), "float32")) |
| 365 | w = relax.Var("w", R.Tensor((dim, dim), "float32")) |
| 366 | with bb.function("main", [x, w]): |
| 367 | with bb.dataflow(): |
| 368 | lv0 = bb.emit(relax.Call(fused_split, (x,))) |
| 369 | gv = bb.emit_output(bb.call_te(topi.nn.dense, lv0, w)) |
| 370 | bb.emit_func_output(gv) |
| 371 | |
| 372 | return bb.get() |
| 373 | |
| 374 | dim = 10 |
| 375 | _check(before(dim), expected(dim)) |
| 376 | |
| 377 | |
| 378 | def test_tuple_intermediate(): |