mosaic/python/dialects/BUILD

load("//third_party/bazel_rules/rules_python/python:py_library.bzl", "py_library")
load("//third_party/llvm/llvm-project/mlir:tblgen.bzl", "gentbl_filegroup")
load("//third_party/pybind11/google3_utils:build_defs.bzl", "pybind_extension")
load("//tools/build_defs/build_test:build_test.bzl", "build_test")

gentbl_filegroup(
    name = "llo_python_gen_raw",
    tbl_outs = [
        (
            [
                "-gen-python-op-bindings",
                "-bind-dialect=llo",
            ],
            "_llo_gen_raw.py",
        ),
    ],
    tblgen = "//third_party/llvm/llvm-project/mlir:mlir-tblgen",
    td_file = ":llo_python.td",
    deps = [
        "//platforms/xla/mosaic:llo_td_files",
    ],
)

genrule(
    name = "llo_python_gen",
    srcs = ["_llo_gen_raw.py"],
    outs = ["_llo_gen.py"],
    cmd = "cat $(location _llo_gen_raw.py) | sed -e 's/^from \\./from mlir\\.dialects\\./g' > $@",
)

py_library(
    name = "llo_dialect",
    srcs = [
        "_llo_gen.py",
        "_llo_ops_ext.py",
        "llo.py",
    ],
    visibility = [
        "//platforms/xla/mosaic:__subpackages__",
    ],
    deps = [
        ":_llo_ext",
        "//third_party/py/mlir",
    ],
)

pybind_extension(
    name = "_llo_ext",
    srcs = ["llo_ext.cc"],
    deps = [
        "//platforms/xla/mosaic:llo_dialect",
        "//third_party/llvm/llvm-project/mlir:CAPIIR",
        "//third_party/llvm/llvm-project/mlir:IR",
        "//third_party/llvm/llvm-project/mlir:MLIRBindingsPythonHeadersAndDeps",
    ],
)

build_test(
    name = "llo_dialect_build_test",
    targets = [
        ":llo_dialect",
    ],
)