Add python bindings for IREE internal dialects.

I've been running into these recently when wanting to create helper scripts. Also some of these will be helpful for Turbine.

compiler/bindings/python/CMakeLists.txt

@@ -50,6 +50,51 @@ set(CMAKE_PLATFORM_NO_VERSIONED_SONAME 1)
 # Sources
 ################################################################################
 
+declare_mlir_python_sources(IREEPythonSources)
+declare_mlir_python_sources(IREEPythonSources.Dialects
+  ADD_TO_PARENT IREEPythonSources
+)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT IREEPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler"
+  TD_FILE dialects/FlowBinding.td
+  SOURCES dialects/flow.py
+  DIALECT_NAME flow
+)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT IREEPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler"
+  TD_FILE dialects/HALBinding.td
+  SOURCES dialects/hal.py
+  DIALECT_NAME hal
+)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT IREEPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler"
+  TD_FILE dialects/StreamBinding.td
+  SOURCES dialects/stream.py
+  DIALECT_NAME stream
+)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT IREEPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler"
+  TD_FILE dialects/UtilBinding.td
+  SOURCES dialects/util.py
+  DIALECT_NAME util
+)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT IREEPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler"
+  TD_FILE dialects/VMBinding.td
+  SOURCES dialects/vm.py
+  DIALECT_NAME vm
+)
+
 declare_mlir_python_sources(IREECompilerAPIPythonTools
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler"
   SOURCES
@@ -88,6 +133,7 @@ set(_SOURCE_COMPONENTS
   # Local sources.
   IREECompilerAPIPythonTools
   IREECompilerPythonExtensions
+  IREEPythonSources
 
   MLIRPythonSources.Core
 

compiler/bindings/python/iree/compiler/dialects/FlowBinding.td

@@ -0,0 +1,12 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef PYTHON_BINDINGS_FLOW_OPS
+#define PYTHON_BINDINGS_FLOW_OPS
+
+include "iree/compiler/Dialect/Flow/IR/FlowOps.td"
+
+#endif // PYTHON_BINDINGS_FLOW_OPS

compiler/bindings/python/iree/compiler/dialects/HALBinding.td

@@ -0,0 +1,12 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef PYTHON_BINDINGS_HAL_OPS
+#define PYTHON_BINDINGS_HAL_OPS
+
+include "iree/compiler/Dialect/HAL/IR/HALOps.td"
+
+#endif // PYTHON_BINDINGS_HAL_OPS

compiler/bindings/python/iree/compiler/dialects/StreamBinding.td

@@ -0,0 +1,12 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef PYTHON_BINDINGS_STREAM_OPS
+#define PYTHON_BINDINGS_STREAM_OPS
+
+include "iree/compiler/Dialect/Stream/IR/StreamOps.td"
+
+#endif // PYTHON_BINDINGS_STREAM_OPS

compiler/bindings/python/iree/compiler/dialects/UtilBinding.td

@@ -0,0 +1,12 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef PYTHON_BINDINGS_UTIL_OPS
+#define PYTHON_BINDINGS_UTIL_OPS
+
+include "iree/compiler/Dialect/Util/IR/UtilOps.td"
+
+#endif // PYTHON_BINDINGS_UTIL_OPS

compiler/bindings/python/iree/compiler/dialects/VMBinding.td

@@ -0,0 +1,12 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef PYTHON_BINDINGS_VM_OPS
+#define PYTHON_BINDINGS_VM_OPS
+
+include "iree/compiler/Dialect/VM/IR/VMOps.td"
+
+#endif // PYTHON_BINDINGS_VM_OPS

compiler/bindings/python/iree/compiler/dialects/flow.py

@@ -0,0 +1,7 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._flow_ops_gen import *

compiler/bindings/python/iree/compiler/dialects/hal.py

@@ -0,0 +1,7 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._hal_ops_gen import *

compiler/bindings/python/iree/compiler/dialects/stream.py

@@ -0,0 +1,7 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._stream_ops_gen import *

compiler/bindings/python/iree/compiler/dialects/util.py

@@ -0,0 +1,7 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._util_ops_gen import *

compiler/bindings/python/iree/compiler/dialects/vm.py

@@ -0,0 +1,7 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._vm_ops_gen import *

compiler/bindings/python/test/ir/CMakeLists.txt

@@ -11,3 +11,9 @@ iree_py_test(
     "registration_test.py"
 )
 
+iree_py_test(
+  NAME
+    dialects_test
+  SRCS
+    "dialects_test.py"
+)

compiler/bindings/python/test/ir/dialects_test.py

@@ -0,0 +1,16 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from iree.compiler import ir
+
+# Make sure that our dialects import.
+from iree.compiler.dialects import (
+    flow,
+    hal,
+    stream,
+    vm,
+    util,
+)