Method entry_a

tests/python/relax/test_transform_allocate_workspace.py:123–137

(
        q: R.Tensor((32, 8, 16, 8), dtype="float16"),
        k: R.Tensor((32, 8, 16, 8), dtype="float16"),
        v: R.Tensor((32, 8, 16, 8), dtype="float16"),
    )

    @R.function
    def entry_a(
        q: R.Tensor((32, 8, 16, 8), dtype="float16"),
        k: R.Tensor((32, 8, 16, 8), dtype="float16"),
        v: R.Tensor((32, 8, 16, 8), dtype="float16"),
    ) -> R.Tensor((32, 8, 16, 8), dtype="float16"):
        cls = Expected
        with R.dataflow():
            workspace_main: R.Tensor((65536,), dtype="uint8") = R.builtin.alloc_tensor(
                R.shape([65536]), R.dtype("uint8"), R.prim_value(0)
            )
            gv: R.Tensor((32, 8, 16, 8), dtype="float16") = cls.fused_relax_nn_attention_cutlass1(
                q, k, v, workspace_main
            )
            R.output(gv)
        return gv

    @R.function
    def entry_b(