Function main

python/tvm/exec/gpu_memory_bandwidth.py:178–258

Entry point

()

def main():  # pylint: disable=too-many-locals
    """Entry point"""
    args = _parse_args()
    # pylint: disable=invalid-name
    target = tvm.target.Target(args.target)
    if args.target_host is not None:
        target = tvm.target.Target(args.target, host=args.target_host)
    dtype = args.dtype
    rpcConfig = None
    if args.rpc_host is not None and args.rpc_port is not None and args.rpc_key is not None:
        rpcConfig = RPCConfig(
            tracker_host=args.rpc_host,
            tracker_port=args.rpc_port,
            tracker_key=args.rpc_key,
            session_priority=1,
            session_timeout_sec=600,
        )

    a = np.random.uniform(-1, 1, (args.xo, args.k, args.xi)).astype(dtype)
    b = np.zeros((args.xo, args.xi), dtype=dtype)
    num_bytes = a.size * a.itemsize + b.size * b.itemsize
    print("###### Bandwidth Test ######")
    print(
        f"Workload [XO, K, XI] => [XO, XI]. "
        f"[{args.xo}, {args.k}, {args.xi}] => [{args.xo}, {args.xi}]"
    )
    print(f"Input size: {num_bytes / 1048576} MB")
    print(f"Target: {target}")

    # pylint: enable=invalid-name
    best_bandwidth = -1
    for len_bx, len_tx, len_vec in itertools.product(
        args.bx,
        args.tx,
        args.vec,
    ):
        func = _workload(
            len_xo=args.xo,
            len_k=args.k,
            len_xi=args.xi,
            dtype=dtype,
        )
        sch = tvm.s_tir.Schedule(func)
        _schedule(sch, len_bx, len_tx, len_vec)

        if rpcConfig is None:
            _, profile_result = local_run(
                tvm.compile(sch.mod, target=target),
                target.kind.name,
                [a, b],
                evaluator_config=EvaluatorConfig(
                    number=10,
                    repeat=1,
                    min_repeat_ms=100,
                    enable_cpu_cache_flush=False,
                ),
            )
        else: