From 6864d27088e6aa4c6420511b185b7cffc49d348a Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko@google.com>
Date: Wed, 16 Mar 2022 15:59:59 +0100
Subject: [PATCH] [mlir][bazel] make .pyi files available to Bazel

These files are necessary for various type checking and autocompletion
tooling to work.

Differential Revision: https://reviews.llvm.org/D121810
---
 .../mlir/python/BUILD.bazel                   | 87 ++++++++++++++++++-
 1 file changed, 84 insertions(+), 3 deletions(-)

diff --git a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
index 0dce73a4fa2a6a..f322f3df7705c7 100644
--- a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
@@ -42,6 +42,13 @@ filegroup(
     ],
 )
 
+filegroup(
+    name = "ExecutionEnginePyIFiles",
+    srcs = [
+        "mlir/_mlir_libs/_mlirExecutionEngine.pyi",
+    ],
+)
+
 filegroup(
     name = "IRPyFiles",
     srcs = [
@@ -49,6 +56,14 @@ filegroup(
     ],
 )
 
+filegroup(
+    name = "IRPyIFiles",
+    srcs = [
+        "mlir/_mlir_libs/_mlir/__init__.pyi",
+        "mlir/_mlir_libs/_mlir/ir.pyi",
+    ],
+)
+
 filegroup(
     name = "PassManagerPyFiles",
     srcs = [
@@ -56,6 +71,13 @@ filegroup(
     ],
 )
 
+filegroup(
+    name = "PassManagerPyIFiles",
+    srcs = [
+        "mlir/_mlir_libs/_mlir/passmanager.pyi",
+    ],
+)
+
 filegroup(
     name = "RuntimePyFiles",
     srcs = glob([
@@ -318,16 +340,56 @@ filegroup(
 )
 
 ##---------------------------------------------------------------------------##
-# Quant dialect.
+# PDL dialect.
 ##---------------------------------------------------------------------------##
 
+td_library(
+    name = "PDLPyTdFiles",
+    srcs = [
+        "//mlir:include/mlir/Bindings/Python/Attributes.td",
+    ],
+    includes = ["../include"],
+    deps = [
+        "//mlir:OpBaseTdFiles",
+        "//mlir:PDLDialectTdFiles",
+    ],
+)
+
+gentbl_filegroup(
+    name = "PDLPyGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-python-op-bindings",
+                "-bind-dialect=pdl",
+            ],
+            "mlir/dialects/_pdl_ops_gen.py",
+        ),
+    ],
+    tblgen = "//mlir:mlir-tblgen",
+    td_file = "mlir/dialects/PDLOps.td",
+    deps = [
+        ":PDLPyTdFiles",
+    ],
+)
+
 filegroup(
-    name = "QuantPyFiles",
+    name = "PDLPyIFiles",
     srcs = [
-        "mlir/dialects/quant.py",
+        "mlir/_mlir_libs/_mlir/dialects/pdl.pyi",
+    ],
+)
+
+filegroup(
+    name = "PDLPyFiles",
+    srcs = [
+        "mlir/dialects/_pdl_ops_ext.py",
+        "mlir/dialects/pdl.py",
+        ":PDLPyGen",
     ],
 )
 
+
 ##---------------------------------------------------------------------------##
 # PythonTest dialect.
 ##---------------------------------------------------------------------------##
@@ -369,6 +431,25 @@ filegroup(
     ],
 )
 
+##---------------------------------------------------------------------------##
+# Quant dialect.
+##---------------------------------------------------------------------------##
+
+filegroup(
+    name = "QuantPyIFiles",
+    srcs = [
+        "mlir/_mlir_libs/_mlir/dialects/quant.pyi",
+    ],
+)
+
+filegroup(
+    name = "QuantPyFiles",
+    srcs = [
+        "mlir/dialects/quant.py",
+    ],
+)
+
+
 ##---------------------------------------------------------------------------##
 # SCF dialect.
 ##---------------------------------------------------------------------------##