MLIR bazel build and test

.github/workflows/enzyme-bazel.yml

@@ -48,4 +48,5 @@ jobs:
     - name: cmake
       run: |
           cd enzyme
-          "${GITHUB_WORKSPACE}/bin/bazel" build :EnzymeStatic 
+          "${GITHUB_WORKSPACE}/bin/bazel" build :EnzymeStatic :enzymemlir-opt
+          "${GITHUB_WORKSPACE}/bin/bazel" test ...

enzyme/BUILD

@@ -1,4 +1,7 @@
 load("@llvm-project//llvm:tblgen.bzl", "gentbl")
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("@llvm-project//llvm:lit_test.bzl", "lit_test", "package_path")
+load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
 
 licenses(["notice"])
 
@@ -213,3 +216,251 @@ genrule(
     cmd = "cp $< $@",
     output_to_bindir = 1,
 )
+
+td_library(
+    name = "EnzymeDialectTdFiles",
+    srcs = [
+    	"Enzyme/MLIR/Dialect/Dialect.td",
+    ],
+	deps = [
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+        "@llvm-project//mlir:ViewLikeInterfaceTdFiles",
+        "@llvm-project//mlir:FunctionInterfacesTdFiles",
+        "@llvm-project//mlir:ControlFlowInterfacesTdFiles",
+        "@llvm-project//mlir:LoopLikeInterfaceTdFiles",
+	]
+)
+
+gentbl_cc_library(
+    name = "EnzymeOpsIncGen",
+    tbl_outs = [
+        (
+            ["-gen-op-decls"],
+            "Enzyme/MLIR/Dialect/EnzymeOps.h.inc",
+        ),
+        (
+            ["-gen-op-defs"],
+            "Enzyme/MLIR/Dialect/EnzymeOps.cpp.inc",
+        ),
+        (
+            [
+                "-gen-dialect-decls",
+                "-dialect=enzyme",
+            ],
+            "Enzyme/MLIR/Dialect/EnzymeOpsDialect.h.inc",
+        ),
+        (
+            [
+                "-gen-dialect-defs",
+                "-dialect=enzyme",
+            ],
+            "Enzyme/MLIR/Dialect/EnzymeOpsDialect.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Enzyme/MLIR/Dialect/EnzymeOps.td",
+    deps = [":EnzymeDialectTdFiles"],
+)
+
+td_library(
+    name = "EnzymePassesTdFiles",
+    srcs = [
+    ],
+	deps = [
+        "@llvm-project//mlir:PassBaseTdFiles",
+	]
+)
+
+gentbl_cc_library(
+    name = "EnzymePassesIncGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-pass-decls",
+                "-name=enzyme",
+            ],
+            "Enzyme/MLIR/Passes/Passes.h.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Enzyme/MLIR/Passes/Passes.td",
+    deps = [":EnzymePassesTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "EnzymeTypesIncGen",
+    tbl_outs = [
+        (
+            ["-gen-typedef-decls"],
+            "Enzyme/MLIR/Dialect/EnzymeOpsTypes.h.inc",
+        ),
+        (
+            ["-gen-typedef-defs"],
+            "Enzyme/MLIR/Dialect/EnzymeOpsTypes.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Enzyme/MLIR/Dialect/EnzymeOps.td",
+    deps = [":EnzymeDialectTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "EnzymeEnumsIncGen",
+    tbl_outs = [
+        (
+            ["-gen-enum-decls"],
+            "Enzyme/MLIR/Dialect/EnzymeEnums.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "Enzyme/MLIR/Dialect/EnzymeEnums.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Enzyme/MLIR/Dialect/EnzymeOps.td",
+    deps = [":EnzymeDialectTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "EnzymeAttributesIncGen",
+    tbl_outs = [
+        (
+            ["-gen-attrdef-decls"],
+            "Enzyme/MLIR/Dialect/EnzymeAttributes.h.inc",
+        ),
+        (
+            ["-gen-attrdef-defs"],
+            "Enzyme/MLIR/Dialect/EnzymeAttributes.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Enzyme/MLIR/Dialect/EnzymeOps.td",
+    deps = [":EnzymeDialectTdFiles"],
+)
+
+
+gentbl_cc_library(
+    name = "EnzymeTypeInterfacesIncGen",
+    tbl_outs = [
+        (
+            ["--gen-type-interface-decls"],
+            "Enzyme/MLIR/Interfaces/AutoDiffTypeInterface.h.inc",
+        ),
+        (
+            ["--gen-type-interface-defs"],
+            "Enzyme/MLIR/Interfaces/AutoDiffTypeInterface.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Enzyme/MLIR/Interfaces/AutoDiffTypeInterface.td",
+    deps = [":EnzymeDialectTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "EnzymeOpInterfacesIncGen",
+    tbl_outs = [
+        (
+            ["--gen-op-interface-decls"],
+            "Enzyme/MLIR/Interfaces/AutoDiffOpInterface.h.inc",
+        ),
+        (
+            ["--gen-op-interface-defs"],
+            "Enzyme/MLIR/Interfaces/AutoDiffOpInterface.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Enzyme/MLIR/Interfaces/AutoDiffOpInterface.td",
+    deps = [":EnzymeDialectTdFiles"],
+)
+
+cc_library(
+    name = "EnzymeMLIR",
+    srcs = glob([
+        "Enzyme/MLIR/Dialect/*.cpp",
+        "Enzyme/MLIR/Passes/*.cpp",
+        "Enzyme/MLIR/Interfaces/*.cpp",
+        "Enzyme/MLIR/Analysis/*.cpp",
+        "Enzyme/MLIR/Implementations/*.cpp",
+    ]),
+    hdrs = glob([
+        "Enzyme/MLIR/Dialect/*.h",
+        "Enzyme/MLIR/Passes/*.h",
+        "Enzyme/MLIR/Interfaces/*.h",
+        "Enzyme/MLIR/Analysis/*.h",
+        "Enzyme/MLIR/Implementations/*.h",
+        "Enzyme/Utils.h",
+        "Enzyme/TypeAnalysis/*.h"
+    ]),
+    includes = ["Enzyme/MLIR", "Enzyme"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":EnzymeOpsIncGen",
+        ":EnzymePassesIncGen",
+		":EnzymeTypesIncGen",
+		":EnzymeEnumsIncGen",
+		":EnzymeAttributesIncGen",
+		":EnzymeTypeInterfacesIncGen",
+		":EnzymeOpInterfacesIncGen",
+        "@llvm-project//mlir:AffineDialect",
+        "@llvm-project//mlir:LLVMCommonConversion",
+        "@llvm-project//mlir:ConversionPasses",
+        "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:GPUDialect",
+        "@llvm-project//mlir:AsyncDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncExtensions",
+        "@llvm-project//mlir:MathDialect",
+        "@llvm-project//mlir:MemRefDialect",
+    ],
+)
+
+cc_binary(
+    name = "enzymemlir-opt",
+    srcs = ["Enzyme/MLIR/enzymemlir-opt.cpp"],
+    visibility = ["//visibility:public"],
+    includes = ["Enzyme/MLIR"],
+    deps = [
+        ":EnzymeMLIR",
+        "@llvm-project//mlir:MlirOptLib",
+        "@llvm-project//mlir:AllPassesAndDialects",
+    ],
+)
+
+# Generates lit config input file by applying path placeholder substitutions
+# similar to the configure_lit_site_cfg CMake macro.
+expand_template(
+    name = "lit_site_cfg_py",
+    testonly = True,
+    out = "test/lit.site.cfg.py",
+    substitutions = {
+        "@LLVM_VERSION_MAJOR@": "18",
+        "@LIT_SITE_CFG_IN_HEADER@": "# Autogenerated, do not edit.",
+        "@LLVM_BINARY_DIR@": package_path("@llvm-project//llvm:BUILD"),
+        "@LLVM_TOOLS_BINARY_DIR@": package_path("@llvm-project//llvm:BUILD"),
+        "@LLVM_LIBS_DIR@": package_path("@llvm-project//llvm:BUILD"),
+        "@ENZYME_SOURCE_DIR@": "",
+        "@ENZYME_BINARY_DIR@": "",
+        "@TARGET_TRIPLE@": "",
+        "@TARGETS_TO_BUILD@": "ALL",
+        "@LLVM_SHLIBEXT@": ".so",
+    },
+    template = "test/lit.site.cfg.py.in",
+    visibility = ["//visibility:private"],
+)
+
+[
+    lit_test(
+        name = "%s.test" % src,
+        srcs = [src],
+        data = [
+            ":test/lit.cfg.py",
+            ":test/lit.site.cfg.py",
+        	"@llvm-project//llvm:FileCheck",
+        	"@llvm-project//llvm:count",
+        	"@llvm-project//llvm:not",
+			":enzymemlir-opt"
+        ],
+    )
+    for src in glob(["test/**/*.mlir"])
+]

enzyme/WORKSPACE

@@ -7,7 +7,7 @@ new_local_repository(
 
 load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure")
 
-llvm_configure(name = "llvm-project", targets = ["X86"])
+llvm_configure(name = "llvm-project", targets = ["X86", "NVPTX"])
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")
@@ -44,3 +44,19 @@ maybe(
         "https://github.com/facebook/zstd/releases/download/v1.5.2/zstd-1.5.2.tar.gz"
     ],
 )
+
+PYRULES_COMMIT = "fe33a4582c37499f3caeb49a07a78fc7948a8949"
+PYRULES_SHA256 = "cfa6957832ae0e0c7ee2ccf455a888a291e8419ed8faf45f4420dd7414d5dd96"
+
+http_archive(
+    name = "rules_python",
+    sha256 = PYRULES_SHA256,
+    strip_prefix = "rules_python-" + PYRULES_COMMIT,
+    urls = ["https://github.com/bazelbuild/rules_python/archive/{commit}.tar.gz".format(commit = PYRULES_COMMIT)]
+)
+
+
+load("@rules_python//python:repositories.bzl", "py_repositories")
+
+py_repositories()
+

enzyme/test/lit.cfg.py

@@ -45,6 +45,7 @@
 #llvm_config.add_tool_substitutions(tools, config.llvm_tools_dir)
 
 # opt knows whether it is compiled with -DNDEBUG.
+"""
 import subprocess
 try:
     opt_cmd = subprocess.Popen([os.path.join(config.llvm_tools_dir, 'opt'), '-version'],
@@ -68,3 +69,4 @@
 except OSError:
     print("Could not find llvm-config in " + config.llvm_tools_dir)
     exit(42)
+"""

enzyme/test/lit.site.cfg.py.in

@@ -1,5 +1,7 @@
 @LIT_SITE_CFG_IN_HEADER@
 
+import os
+
 #config.llvm_src_root = "@LLVM_SOURCE_DIR@"
 config.llvm_ver = "@LLVM_VERSION_MAJOR@"
 config.llvm_obj_root = "@LLVM_BINARY_DIR@"
@@ -40,7 +42,12 @@ config.substitutions.append(('%shlibext', config.llvm_shlib_ext))
 config.substitutions.append(('%lli', config.llvm_tools_dir + "/lli" + (" --jit-kind=mcjit" if int(config.llvm_ver) >= 13 else "")
 ))
 config.substitutions.append(('%opt', config.llvm_tools_dir + "/opt"))
-config.substitutions.append(('%eopt', config.enzyme_obj_root + "/Enzyme/MLIR/enzymemlir-opt"))
+
+eopt = config.enzyme_obj_root + "/Enzyme/MLIR/enzymemlir-opt"
+if len("@ENZYME_BINARY_DIR@") == 0:
+  eopt = os.path.dirname(os.path.abspath(__file__)) + "/../enzymemlir-opt"
+
+config.substitutions.append(('%eopt', eopt))
 config.substitutions.append(('%llvmver', config.llvm_ver))
 config.substitutions.append(('%FileCheck', config.llvm_tools_dir + "/FileCheck"))
 config.substitutions.append(('%clang', config.llvm_tools_dir + "/clang"))
@@ -81,4 +88,7 @@ config.substitutions.append(('%loadClangEnzyme', oldPM if int(config.llvm_ver) <
 config.substitutions.append(('%newLoadClangEnzyme', newPM))
 
 # Let the main config do the real work.
-lit_config.load_config(config, "@ENZYME_SOURCE_DIR@/test/lit.cfg.py")
+cfgfile = "@ENZYME_SOURCE_DIR@/test/lit.cfg.py"
+if len("@ENZYME_SOURCE_DIR@") == 0:
+  cfgfile = os.path.dirname(os.path.abspath(__file__)) + "/lit.cfg.py"
+lit_config.load_config(config, cfgfile)