Function _workload

python/tvm/exec/gpu_memory_bandwidth.py:138–153

(
    len_xo: int,
    len_k: int,
    len_xi: int,
    dtype: str,
)

def _workload(
    len_xo: int,
    len_k: int,
    len_xi: int,
    dtype: str,
):
    # pylint: disable=invalid-name
    A = te.placeholder((len_xo, len_k, len_xi), dtype=dtype, name="A")
    k = te.reduce_axis((0, len_k), "k")
    B = te.compute(
        (len_xo, len_xi),
        lambda i, j: te.sum(A[i, k, j], axis=k),
        name="B",
    )
    # pylint: enable=invalid-name
    return te.create_prim_func([A, B])


def _schedule(