Method main

tests/python/relax/test_transform_codegen_pass.py:341–363

(
            x: R.Tensor((1, 4096), dtype="float16"),
            w1: R.Tensor((4096, "r1"), dtype="float16"),
            w2: R.Tensor((4096, "r2"), dtype="float16"),
        )

    class Expected:
        @R.function
        def main(
            x: R.Tensor((1, 4096), dtype="float16"),
            w1: R.Tensor((4096, "r1"), dtype="float16"),
            w2: R.Tensor((4096, "r2"), dtype="float16"),
        ) -> R.Tuple(R.Tensor((1, "r1"), dtype="float16"), R.Tensor((1, "r2"), dtype="float16")):
            r1 = T.int64()
            r2 = T.int64()
            with R.dataflow():
                lv = R.call_dps_packed(
                    "fused_relax_matmul_cublas",
                    (x, w1),
                    out_sinfo=R.Tensor((1, r1), dtype="float16"),
                )
                lv1 = R.call_dps_packed(
                    "fused_relax_matmul_cublas",
                    (x, w2),
                    out_sinfo=R.Tensor((1, r2), dtype="float16"),
                )
                gv: R.Tuple(
                    R.Tensor((1, r1), dtype="float16"), R.Tensor((1, r2), dtype="float16")
                ) = (lv, lv1)
                R.output(gv)
            return gv

    after = relax.transform.RunCodegen()(Before)
    tvm.ir.assert_structural_equal(after["main"], Expected["main"])