from generator import (
    GenerateGemmOperations,
    GenerateGemvOperations,
    GenerateConv2dOperations,
    GenerateDeconvOperations,
    GenerateDwconv2dFpropOperations,
    GenerateDwconv2dDgradOperations,
    GenerateDwconv2dWgradOperations,
)


class GenArg:
    def __init__(self, gen_op, gen_type):
        self.operations = gen_op
        self.type = gen_type


def write_op_list(f, gen_op, gen_type):
    if gen_op == "gemm":
        operations = GenerateGemmOperations(GenArg(gen_op, gen_type))
    elif gen_op == "gemv":
        operations = GenerateGemvOperations(GenArg(gen_op, gen_type))
    elif gen_op == "conv2d":
        operations = GenerateConv2dOperations(GenArg(gen_op, gen_type))
    elif gen_op == "deconv":
        operations = GenerateDeconvOperations(GenArg(gen_op, gen_type))
    elif gen_op == "dwconv2d_fprop":
        operations = GenerateDwconv2dFpropOperations(GenArg(gen_op, gen_type))
    elif gen_op == "dwconv2d_dgrad":
        operations = GenerateDwconv2dDgradOperations(GenArg(gen_op, gen_type))
    elif gen_op == "dwconv2d_wgrad":
        operations = GenerateDwconv2dWgradOperations(GenArg(gen_op, gen_type))
    for op in operations:
        f.write('    "%s.cu",\n' % op.procedural_name())
    if gen_op != "gemv":
        f.write('    "all_%s_%s_operations.cu",\n' % (gen_op, gen_type))


if __name__ == "__main__":
    with open("list.bzl", "w") as f:
        f.write("# Generated by dnn/scripts/cutlass_generator/gen_list.py\n\n")
        f.write("cutlass_gen_list = [\n")
        write_op_list(f, "gemm", "simt")
        write_op_list(f, "gemm", "tensorop1688")
        write_op_list(f, "gemm", "tensorop884")
        write_op_list(f, "gemv", "simt")
        write_op_list(f, "deconv", "simt")
        write_op_list(f, "deconv", "tensorop8816")
        write_op_list(f, "conv2d", "simt")
        write_op_list(f, "conv2d", "tensorop8816")
        write_op_list(f, "conv2d", "tensorop8832")
        write_op_list(f, "dwconv2d_fprop", "simt")
        write_op_list(f, "dwconv2d_fprop", "tensorop884")
        write_op_list(f, "dwconv2d_dgrad", "simt")
        write_op_list(f, "dwconv2d_dgrad", "tensorop884")
        write_op_list(f, "dwconv2d_wgrad", "simt")
        write_op_list(f, "dwconv2d_wgrad", "tensorop884")
        f.write("]")