| 193 | |
| 194 | |
| 195 | def test_recursive(): |
| 196 | scalar_struct_info = rx.TensorStructInfo(shape=[], dtype="int32") |
| 197 | gv0 = rx.Var("gv0", scalar_struct_info) |
| 198 | f = rx.Var("f", rx.FuncStructInfo([scalar_struct_info], scalar_struct_info)) |
| 199 | ipt = rx.Var("ipt", scalar_struct_info) |
| 200 | x0 = rx.Var("x0", scalar_struct_info) |
| 201 | x1 = rx.Var("x1", scalar_struct_info) |
| 202 | x2 = rx.Var("x2", scalar_struct_info) |
| 203 | y = rx.Var("y", scalar_struct_info) |
| 204 | inner_block = rx.BindingBlock( |
| 205 | [rx.VarBinding(x0, rx.const(2, "int32")), rx.VarBinding(y, rx.Call(f, [x0]))] |
| 206 | ) |
| 207 | inner_func = rx.Function([ipt], rx.SeqExpr([inner_block], y), scalar_struct_info) |
| 208 | outer_block = rx.BindingBlock( |
| 209 | [ |
| 210 | rx.VarBinding(f, inner_func), |
| 211 | rx.VarBinding(x1, rx.const(1, "int32")), |
| 212 | rx.VarBinding(x2, rx.op.add(x1, rx.Call(f, [x1]))), |
| 213 | rx.VarBinding(gv0, x2), |
| 214 | ] |
| 215 | ) |
| 216 | func = rx.Function([], rx.SeqExpr([outer_block], gv0), scalar_struct_info) |
| 217 | mod = tvm.IRModule.from_expr(func) |
| 218 | normalized = rx.transform.Normalize()(mod) |
| 219 | rx.analysis.well_formed(normalized) |
| 220 | |
| 221 | |
| 222 | def test_if(): |