[llvm] [bazel] Add a bazel flag to enable building MLIR with CUDA support (PR #88856)

[Adam Paszke via llvm-commits]

https://github.com/apaszke created https://github.com/llvm/llvm-project/pull/88856

This makes it possible to specify `-- at llvm-project//mlir:enable_cuda=true` on the bazel command line and get a build that includes NVIDIA GPU support in MLIR.

>From c1ce6a16f1ddd8f8fb3d29ef67c68eaac4e8e2f7 Mon Sep 17 00:00:00 2001
From: Adam Paszke <adam.paszke at gmail.com>
Date: Mon, 15 Apr 2024 15:56:10 +0000
Subject: [PATCH] [bazel] Add a bazel flag to enable building MLIR with CUDA
 support

---
 utils/bazel/llvm-project-overlay/mlir/BUILD.bazel   | 13 +++++++++++++
 .../bazel/llvm-project-overlay/mlir/build_defs.bzl  |  3 ++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 67052fcd399309..d268f4046bc192 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -6,6 +6,7 @@
 #   The MLIR "Multi-Level Intermediate Representation" Compiler Infrastructure
 
 load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
+load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
 load(
     ":build_defs.bzl",
     "cc_headers_only",
@@ -28,6 +29,18 @@ exports_files([
     "utils/textmate/mlir.json",
 ])
 
+bool_flag(
+    name = "enable_cuda",
+    build_setting_default = False,
+)
+
+config_setting(
+    name = "enable_cuda_config",
+    flag_values = {
+        ":enable_cuda": "True",
+    },
+)
+
 expand_template(
     name = "mlir_config_h_gen",
     out = "include/mlir/Config/mlir-config.h",
diff --git a/utils/bazel/llvm-project-overlay/mlir/build_defs.bzl b/utils/bazel/llvm-project-overlay/mlir/build_defs.bzl
index fe39dbe21e3390..53767fd03c44ad 100644
--- a/utils/bazel/llvm-project-overlay/mlir/build_defs.bzl
+++ b/utils/bazel/llvm-project-overlay/mlir/build_defs.bzl
@@ -6,7 +6,8 @@
 
 def if_cuda_available(if_true, if_false = []):
     return select({
-        # CUDA is not yet supported.
+        # CUDA auto-detection is not yet supported.
+        "//mlir:enable_cuda_config": if_true,
         "//conditions:default": if_false,
     })