| 269 | |
| 270 | |
| 271 | def test_fuse_tuple_get_elemwise(): |
| 272 | def before(dim: int): |
| 273 | bb = relax.BlockBuilder() |
| 274 | x = relax.Var("x", R.Tensor((1, dim), "float32")) |
| 275 | w = relax.Var("w", R.Tensor((3 * dim, dim), "float32")) |
| 276 | with bb.function("main", [x, w]): |
| 277 | with bb.dataflow(): |
| 278 | lv0 = bb.emit_te(topi.nn.dense, x, w) |
| 279 | lv1 = bb.emit_te(topi.split, lv0, indices_or_sections=3, axis=1) |
| 280 | lv2 = bb.emit(relax.TupleGetItem(lv1, 0)) |
| 281 | lv3 = bb.emit_te(topi.sigmoid, lv2) |
| 282 | lv4 = bb.emit(relax.TupleGetItem(lv1, 1)) |
| 283 | lv5 = bb.emit_te(topi.tanh, lv4) |
| 284 | lv6 = bb.emit(relax.TupleGetItem(lv1, 2)) |
| 285 | lv7 = bb.emit_te(topi.exp, lv6) |
| 286 | lv8 = bb.emit_te(topi.multiply, lv5, lv7) |
| 287 | gv = bb.emit_output(bb.call_te(topi.add, lv3, lv8)) |
| 288 | bb.emit_func_output(gv) |
| 289 | |
| 290 | return bb.get() |
| 291 | |
| 292 | def expected(dim: int): |
| 293 | bb = relax.BlockBuilder() |
| 294 | |
| 295 | # Grouped function |
| 296 | dense = relax.Var("dense", R.Tensor((1, 3 * dim), "float32")) |
| 297 | with bb.function( |
| 298 | "fused_split_sigmoid_tanh_exp_multiply_add", |
| 299 | [dense], |
| 300 | attrs={"Primitive": True}, |
| 301 | private=True, |
| 302 | ): |
| 303 | with bb.dataflow(): |
| 304 | lv0 = bb.emit_te(topi.split, dense, indices_or_sections=3, axis=1) |
| 305 | lv1 = bb.emit(relax.TupleGetItem(lv0, 0)) |
| 306 | lv2 = bb.emit_te(topi.sigmoid, lv1) |
| 307 | lv3 = bb.emit(relax.TupleGetItem(lv0, 1)) |
| 308 | lv4 = bb.emit_te(topi.tanh, lv3) |
| 309 | lv5 = bb.emit(relax.TupleGetItem(lv0, 2)) |
| 310 | lv6 = bb.emit_te(topi.exp, lv5) |
| 311 | lv7 = bb.emit_te(topi.multiply, lv4, lv6) |
| 312 | gv = bb.emit_output(bb.call_te(topi.add, lv2, lv7)) |
| 313 | bb.emit_func_output(gv) |
| 314 | |
| 315 | # Get the global variables of the grouped functions |
| 316 | fused_split_sigmoid_tanh_exp_multiply_add = bb.get().get_global_var( |
| 317 | "fused_split_sigmoid_tanh_exp_multiply_add" |
| 318 | ) |
| 319 | |
| 320 | # Main function |
| 321 | x = relax.Var("x", R.Tensor((1, dim), "float32")) |
| 322 | w = relax.Var("w", R.Tensor((3 * dim, dim), "float32")) |
| 323 | with bb.function("main", [x, w]): |
| 324 | with bb.dataflow(): |
| 325 | lv0 = bb.emit_te(topi.nn.dense, x, w) |
| 326 | gv = bb.emit_output(relax.Call(fused_split_sigmoid_tanh_exp_multiply_add, (lv0,))) |
| 327 | bb.emit_func_output(gv) |
| 328 | |