# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# gemm tuple shape (M, N, K)
gemm_A = [
    (128, 4096, 4096),
    (256, 4096, 4096),
    (512, 4096, 4096),
    (1024, 4096, 4096),
    (128, 1024, 1024),
    (256, 1024, 1024),
    (512, 1024, 1024),
    (1024, 1024, 1024),
    (4096, 4096, 128),
    (4096, 4096, 256),
    (4096, 4096, 512),
    (4096, 4096, 1024),
    (1024, 1024, 128),
    (1024, 1024, 256),
    (1024, 1024, 512),
]

gemm_B = [
    (128, 4096, 40928),
    (256, 4096, 40928),
    (512, 4096, 40928),
    (1024, 4096, 40928),
    (128, 40928, 4096),
    (256, 40928, 4096),
    (512, 40928, 4096),
    (1024, 40928, 4096),
    (128, 1024, 2000),
    (256, 1024, 2000),
    (512, 1024, 2000),
    (1024, 1024, 2000),
    (1024, 2000, 128),
    (1024, 2000, 256),
    (1024, 2000, 512),
    (1024, 2000, 1024),
    (4096, 40928, 128),
    (4096, 40928, 256),
    (4096, 40928, 512),
    (4096, 40928, 1024),
]

gemm_C = [
    (1024, 1024, 64),
    (1024, 64, 1024),
    (1024, 4096, 1024),
    (1024, 1024, 4096),
]

# emb tuple (features, embdim, nnz, batch)
emb_A = [
    (14000000, 128, 30, 512),
    (14000000, 128, 30, 1024),
    (14000000, 128, 30, 2048),
    (14000000, 128, 30, 4096),
    (14000000, 128, 30, 8192),
    (14000000, 128, 30, 16384),
    (14000000, 128, 30, 32768),
    (14000000, 128, 30, 65536),
    (26000000, 128, 30, 512),
    (26000000, 128, 30, 1024),
    (26000000, 128, 30, 2048),
    (26000000, 128, 30, 4096),
    (26000000, 128, 30, 8192),
    (26000000, 128, 30, 16384),
    (26000000, 128, 30, 32768),
    (26000000, 128, 30, 65536),
]

emb_B = [
    (4800000, 56, 34, 2048),
    (4800000, 56, 34, 4096),
    (4800000, 56, 34, 8192),
    (4800000, 56, 34, 16384),
    (4800000, 56, 34, 32768),
    (4800000, 56, 34, 65536),
]

# mlp tuple (layer-num, input-size, hidden-size, output-size, batch-size)
mlp_A = [
    (18, 1024, 1024, 1024, 128),
    (18, 1024, 1024, 1024, 256),
    (18, 1024, 1024, 1024, 512),
    (18, 1024, 1024, 1024, 1024),
    (18, 1024, 1024, 1024, 2048),
    (18, 1024, 1024, 1024, 4096),
    (18, 4096, 4096, 4096, 128),
    (18, 4096, 4096, 4096, 256),
    (18, 4096, 4096, 4096, 512),
    (18, 4096, 4096, 4096, 1024),
    (18, 4096, 4096, 4096, 2048),
    (18, 4096, 4096, 4096, 4096),
]
