<commit_before>"""Generates cubin headers for TF dialect ops."""

load("@local_config_cuda//cuda:build_defs.bzl", "cuda_gpu_architectures")
load(
    "@local_config_rocm//rocm:build_defs.bzl",
    "rocm_gpu_architectures",
)
load(
    "//tensorflow/stream_executor:build_defs.bzl",
    "if_gpu_is_configured",
)
load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain")

def _lookup_file(filegroup, path):
    """Extracts file at (relative) path in filegroup."""
    for file in filegroup.files.to_list():
        if file.path.endswith(path) or file.path.endswith(path + ".exe"):
            return file
    return None

GpuBinaryInfo = provider(
    "GPU binaries in either cubin format or hsaco format",
    fields = ["gpu_bins"],
)

type_to_mlir = {
    "c64": "complex<f32>",
    "c128": "complex<f64>",
}

type_to_tf_dtype = {
    "i1": "DT_BOOL",
    "i8": "DT_INT8",
    "i16": "DT_INT16",
    "i32": "DT_INT32",
    "i64": "DT_INT64",
    "ui8": "DT_UINT8",
    "ui16": "DT_UINT16",
    "ui32": "DT_UINT32",
    "ui64": "DT_UINT64",
    "f16": "DT_HALF",
    "f32": "DT_FLOAT",
    "f64": "DT_DOUBLE",
    "c64": "DT_COMPLEX64",
    "c128": "DT_COMPLEX128",
}

def _get_mlir_type(type):
    """Return the mlir type corresponding to 'type'"""
    if type in type_to_mlir:
        return type_to_mlir[type]
    return type

def _gen_mlir_op_impl(ctx):
    mlir_type = _get_mlir_type(ctx.attr.type)
    mlir_output_type = _get_mlir_type(ctx.attr.output_type)

    cmd = ctx.actions.run_shell(
        inputs = [ctx.file.template],
        outputs = [ctx.outputs.out],
        command = (
            ("cat %s | sed 's/platform/%s/g' | sed 's/_elem_type/_%s/g' | " +
             "sed 's/elem_type/%s/g' | " + "sed 's/_output_type/_%s/g' | " +
             "sed 's/output_type/%s/g' > %s") % (
                ctx.file.template.path,
                ctx.attr.platform.upper(),
                type_to_tf_dtype[ctx.attr.type],
                mlir_type,
                type_to_tf_dtype[ctx.attr.output_type],
                mlir_output_type,
                ctx.outputs.out.path,
            )
        ),
        use_default_shell_env = True,
    )

_gen_mlir_op_rule = rule(
    attrs = {
        "template": attr.label(
            mandatory = True,
            allow_single_file = True,
        ),
        "type": attr.string(mandatory = True),
        "output_type": attr.string(mandatory = True),
        "platform": attr.string(mandatory = True),
        "out": attr.output(mandatory = True),
    },
    output_to_genfiles = True,
    implementation = _gen_mlir_op_impl,
)

def _gen_mlir_op(op, type, platform, output_type):
    _gen_mlir_op_rule(
        name = "generate_{op}_{platform}_{type}_{output_type}_mlir".format(
            op = op,
            platform = platform,
            type = type,
            output_type = output_type,
        ),
        out = "{op}_{platform}_{type}_{output_type}.mlir".format(
            op = op,
            platform = platform,
            type = type,
            output_type = output_type,
        ),
        output_type = output_type,
        platform = platform,
        template = "op_definitions/{op}.mlir.tmpl".format(op = op),
        type = type,
    )

################################################################################
# Kernels build rules.
################################################################################

def if_mlir_generated_gpu_kernels_enabled(if_true, if_false = []):
    return select({
        "//tensorflow/core/kernels/mlir_generated:is_gpu_enabled": if_true,
        "//conditions:default": if_false,
    })

def if_mlir_generated_cpu_kernels_enabled(if_true, if_false = []):
    return select({
        "//tensorflow/core/kernels/mlir_generated:is_cpu_enabled": if_true,
        "//conditions:default": if_false,
    })

def if_mlir_generated_experimental_kernels_enabled(if_true, if_false = []):
    return select({
        "//tensorflow/core/kernels/mlir_generated:is_experimental_enabled": if_true,
        "//conditions:default": if_false,
    })

def _gen_kernel_bin_impl(ctx):
    cc_toolchain = find_cpp_toolchain(ctx)
    feature_configuration = cc_common.configure_features(
        ctx = ctx,
        cc_toolchain = cc_toolchain,
        requested_features = ctx.features,
        unsupported_features = ctx.disabled_features,
    )
    cmd_args = list(ctx.attr.extra_args)
    if ctx.attr.unroll_factors:
        cmd_args.append("--unroll_factors=%s" % ctx.attr.unroll_factors)
    gpu_bin = ctx.outputs.kernel

    # cc_binary seems not to bring its dependencies with it, so do that explicitly here.
    ctx.actions.run(
        inputs = [ctx.file.mlir_op, ctx.file._tfso],
        outputs = [gpu_bin],
        executable = ctx.executable._tool,
        arguments = cmd_args + [
            "--tile_sizes=%s" % ctx.attr.tile_size,
            "--max-supported-rank=%s" % ctx.attr.max_supported_rank,
            "--arch=%s" % ",".join(ctx.attr.gpu_archs),
            "--input=%s" % ctx.file.mlir_op.path,
            "--output=%s" % gpu_bin.path,
            "--enable_ftz=%s" % (ctx.attr.data_type == "f32"),
            "--cpu_codegen=%s" % ctx.attr.cpu_codegen,
            "--jit_i64_indexed_for_large_tensors=%s" % ctx.attr.jit_i64_indexed_for_large_tensor,
            "--jit=%s" % ctx.attr.jit,
        ],
        use_default_shell_env = True,
        mnemonic = "compile",
        progress_message = "Generating kernel '%{label}'",
    )
    compilation_outputs = cc_common.create_compilation_outputs(
        # We always produce PIC object files, so use the same object files for both.
        objects = depset([gpu_bin]),
        pic_objects = depset([gpu_bin]),
    )
    (linking_context, linking_outputs) = cc_common.create_linking_context_from_compilation_outputs(
        name = ctx.label.name,<commit_msg>refactor var name<commit_after>"""Generates cubin headers for TF dialect ops."""

load("@local_config_cuda//cuda:build_defs.bzl", "cuda_gpu_architectures")
load(
    "@local_config_rocm//rocm:build_defs.bzl",
    "rocm_gpu_architectures",
)
load(
    "//tensorflow/stream_executor:build_defs.bzl",
    "if_gpu_is_configured",
)
load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain")

def _lookup_file(filegroup, path):
    """Extracts file at (relative) path in filegroup."""
    for file in filegroup.files.to_list():
        if file.path.endswith(path) or file.path.endswith(path + ".exe"):
            return file
    return None

GpuBinaryInfo = provider(
    "GPU binaries in either cubin format or hsaco format",
    fields = ["gpu_bins"],
)

type_to_mlir = {
    "c64": "complex<f32>",
    "c128": "complex<f64>",
}

type_to_tf_dtype = {
    "i1": "DT_BOOL",
    "i8": "DT_INT8",
    "i16": "DT_INT16",
    "i32": "DT_INT32",
    "i64": "DT_INT64",
    "ui8": "DT_UINT8",
    "ui16": "DT_UINT16",
    "ui32": "DT_UINT32",
    "ui64": "DT_UINT64",
    "f16": "DT_HALF",
    "f32": "DT_FLOAT",
    "f64": "DT_DOUBLE",
    "c64": "DT_COMPLEX64",
    "c128": "DT_COMPLEX128",
}

def _get_mlir_type(type):
    """Return the mlir type corresponding to 'type'"""
    if type in type_to_mlir:
        return type_to_mlir[type]
    return type

def _gen_mlir_op_impl(ctx):
    mlir_type = _get_mlir_type(ctx.attr.type)
    mlir_output_type = _get_mlir_type(ctx.attr.output_type)

    cmd = ctx.actions.run_shell(
        inputs = [ctx.file.template],
        outputs = [ctx.outputs.out],
        command = (
            ("cat %s | sed 's/platform/%s/g' | sed 's/_elem_type/_%s/g' | " +
             "sed 's/elem_type/%s/g' | " + "sed 's/_output_type/_%s/g' | " +
             "sed 's/output_type/%s/g' > %s") % (
                ctx.file.template.path,
                ctx.attr.platform.upper(),
                type_to_tf_dtype[ctx.attr.type],
                mlir_type,
                type_to_tf_dtype[ctx.attr.output_type],
                mlir_output_type,
                ctx.outputs.out.path,
            )
        ),
        use_default_shell_env = True,
    )

_gen_mlir_op_rule = rule(
    attrs = {
        "template": attr.label(
            mandatory = True,
            allow_single_file = True,
        ),
        "type": attr.string(mandatory = True),
        "output_type": attr.string(mandatory = True),
        "platform": attr.string(mandatory = True),
        "out": attr.output(mandatory = True),
    },
    output_to_genfiles = True,
    implementation = _gen_mlir_op_impl,
)

def _gen_mlir_op(op, type, platform, output_type):
    _gen_mlir_op_rule(
        name = "generate_{op}_{platform}_{type}_{output_type}_mlir".format(
            op = op,
            platform = platform,
            type = type,
            output_type = output_type,
        ),
        out = "{op}_{platform}_{type}_{output_type}.mlir".format(
            op = op,
            platform = platform,
            type = type,
            output_type = output_type,
        ),
        output_type = output_type,
        platform = platform,
        template = "op_definitions/{op}.mlir.tmpl".format(op = op),
        type = type,
    )

################################################################################
# Kernels build rules.
################################################################################

def if_mlir_generated_gpu_kernels_enabled(if_true, if_false = []):
    return select({
        "//tensorflow/core/kernels/mlir_generated:is_gpu_enabled": if_true,
        "//conditions:default": if_false,
    })

def if_mlir_generated_cpu_kernels_enabled(if_true, if_false = []):
    return select({
        "//tensorflow/core/kernels/mlir_generated:is_cpu_enabled": if_true,
        "//conditions:default": if_false,
    })

def if_mlir_generated_experimental_kernels_enabled(if_true, if_false = []):
    return select({
        "//tensorflow/core/kernels/mlir_generated:is_experimental_enabled": if_true,
        "//conditions:default": if_false,
    })

def _gen_kernel_bin_impl(ctx):
    cc_toolchain = find_cpp_toolchain(ctx)
    feature_configuration = cc_common.configure_features(
        ctx = ctx,
        cc_toolchain = cc_toolchain,
        requested_features = ctx.features,
        unsupported_features = ctx.disabled_features,
    )
    cmd_args = list(ctx.attr.extra_args)
    if ctx.attr.unroll_factors:
        cmd_args.append("--unroll_factors=%s" % ctx.attr.unroll_factors)
    gpu_bin = ctx.outputs.kernel

    # cc_binary seems not to bring its dependencies with it, so do that explicitly here.
    ctx.actions.run(
        inputs = [ctx.file.mlir_op, ctx.file._tfso],
        outputs = [gpu_bin],
        executable = ctx.executable._tool,
        arguments = cmd_args + [
            "--tile_sizes=%s" % ctx.attr.tile_size,
            "--max-supported-rank=%s" % ctx.attr.max_supported_rank,
            "--arch=%s" % ",".join(ctx.attr.gpu_archs),
            "--input=%s" % ctx.file.mlir_op.path,
            "--output=%s" % gpu_bin.path,
            "--enable_ftz=%s" % (ctx.attr.data_type == "f32"),
            "--cpu_codegen=%s" % ctx.attr.cpu_codegen,
            "--jit_i64_indexed_for_large_tensors=%s" % ctx.attr.jit_i64_indexed_for_large_tensors,
            "--jit=%s" % ctx.attr.jit,
        ],
        use_default_shell_env = True,
        mnemonic = "compile",
        progress_message = "Generating kernel '%{label}'",
    )
    compilation_outputs = cc_common.create_compilation_outputs(
        # We always produce PIC object files, so use the same object files for both.
        objects = depset([gpu_bin]),
        pic_objects = depset([gpu_bin]),
    )
    (linking_context, linking_outputs) = cc_common.create_linking_context_from_compilation_outputs(
        name = ctx.label.name,
{'fraction_non_alphanumeric': 0.15931769722814498, 'fraction_numerical': 0.01023454157782516, 'mean_word_length': 2.6563766760212038, 'pattern_counts': {'":': 54, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7146144', 'n_tokens_mistral': 4057, 'n_tokens_neox': 3958, 'n_words': 655}