feat: build the shardy dialect (#622)

* feat: build the shardy dialect

* fix: remove http_archive

* fix: register dialect

* fix: missing loadDialect

Author: avik-pal

deps/ReactantExtra/API.cpp

@@ -37,6 +37,7 @@
 #include "src/enzyme_ad/jax/Implementations/XLADerivatives.h"
 #include "src/enzyme_ad/jax/Passes/Passes.h"
 #include "llvm/Support/TargetSelect.h"
+#include "shardy/dialect/sdy/ir/dialect.h"
 
 #include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
 #include "stablehlo/dialect/ChloOps.h"
@@ -685,6 +686,7 @@ extern "C" void RegisterDialects(MlirContext cctx) {
   context.loadDialect<mlir::mhlo::MhloDialect>();
   context.loadDialect<mlir::stablehlo::StablehloDialect>();
   context.loadDialect<mlir::chlo::ChloDialect>();
+  context.loadDialect<mlir::sdy::SdyDialect>();
 }
 
 #include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
@@ -713,6 +715,8 @@ extern "C" void InitializeRegistryAndPasses(MlirDialectRegistry creg) {
   mlir::registerNVVMDialectImport(registry);
   mlir::LLVM::registerInlinerInterface(registry);
 
+  mlir::sdy::registerAllDialects(registry);
+
   /*
     registry.addExtension(+[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
       LLVM::LLVMFunctionType::attachInterface<MemRefInsider>(*ctx);

deps/ReactantExtra/BUILD

@@ -795,6 +795,21 @@ gentbl_cc_library(
     tblgen = "//:mlir-jl-tblgen",
 )
 
+gentbl_cc_library(
+    name = "ShardyJLIncGen",
+    tbl_outs = [(
+            ["--generator=jl-op-defs", "--disable-module-wrap=0"],
+            "Shardy.jl"
+        )
+    ],
+    td_file = "@shardy//shardy/dialect/sdy/ir:ops.td",
+    deps = [
+        "@shardy//shardy/dialect/sdy/ir:sdy_td_files",
+    ],
+    tblgen = "//:mlir-jl-tblgen",
+    includes = ["external/shardy"],
+)
+
 genrule(
     name = "libMLIR_h.jl",
     tags = [

deps/ReactantExtra/make-bindings.jl

@@ -1,8 +1,16 @@
+const bazel_cmd = if !isnothing(Sys.which("bazelisk"))
+    "bazelisk"
+elseif !isnothing(Sys.which("bazel"))
+    "bazel"
+else
+    error("Could not find `bazel` or `bazelisk` in PATH!")
+end
+
 function build_file(output_path)
     file = basename(output_path)
     run(
         Cmd(
-            `bazel build --action_env=JULIA=$(Base.julia_cmd().exec[1]) --action_env=JULIA_DEPOT_PATH=$(Base.DEPOT_PATH) --repo_env HERMETIC_PYTHON_VERSION="3.10" --check_visibility=false --verbose_failures //:$file`;
+            `$(bazel_cmd) build --action_env=JULIA=$(Base.julia_cmd().exec[1]) --action_env=JULIA_DEPOT_PATH=$(Base.DEPOT_PATH) --repo_env HERMETIC_PYTHON_VERSION="3.10" --check_visibility=false --verbose_failures //:$file`;
             dir=@__DIR__,
         ),
     )
@@ -29,6 +37,7 @@ for file in [
     "Affine.jl",
     "TPU.jl",
     "Triton.jl",
+    "Shardy.jl",
 ]
     build_file(joinpath(src_dir, "mlir", "Dialects", file))
 end