Method main

tests/python/relax/test_transform_lower_gpu_ipc_alloc_storage.py:44–56

(shape: R.Shape(["m", "n"]))

    class Expected:
        @R.function(pure=False)
        def main(shape: R.Shape(["m", "n"])):  # type: ignore
            m = T.int64()
            n = T.int64()
            storage: R.Object = R.call_packed(
                "runtime.disco.cuda_ipc.alloc_storage",
                R.shape([m, n]),
                R.dtype("float16"),
                sinfo_args=(R.Object,),
            )
            alloc: R.Tensor((m, n), dtype="float16") = R.memory.alloc_tensor(  # type: ignore
                storage, R.prim_value(0), R.shape([m, n]), R.dtype("float16")
            )
            return alloc

    mod = relax.transform.LowerGPUIPCAllocStorage()(Module)
    tvm.ir.assert_structural_equal(mod, Expected)