[mlir][python] UB dialect python bindings (#157127)

Author: Hardcode84

mlir/include/mlir/Dialect/UB/IR/UBOps.td

@@ -12,7 +12,7 @@
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/AttrTypeBase.td"
 
-include "UBOpsInterfaces.td"
+include "mlir/Dialect/UB/IR/UBOpsInterfaces.td"
 
 def UB_Dialect : Dialect {
   let name = "ub";

mlir/python/CMakeLists.txt

@@ -453,6 +453,14 @@ declare_mlir_dialect_python_bindings(
   DIALECT_NAME tosa
 )
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/UBOps.td
+  SOURCES dialects/ub.py
+  DIALECT_NAME ub
+)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
@@ -847,4 +855,3 @@ add_mlir_python_modules(MLIRPythonModules
   COMMON_CAPI_LINK_LIBS
     MLIRPythonCAPI
 )
-

mlir/python/mlir/dialects/UBOps.td

@@ -0,0 +1,14 @@
+//===-- UBOps.td - Entry point for UB bindings -------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, 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_UB_OPS
+#define PYTHON_BINDINGS_UB_OPS
+
+include "mlir/Dialect/UB/IR/UBOps.td"
+
+#endif // PYTHON_BINDINGS_UB_OPS

mlir/python/mlir/dialects/ub.py

@@ -0,0 +1,5 @@
+#  Part of the LLVM Project, 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 ._ub_ops_gen import *

mlir/test/python/dialects/ub.py

@@ -0,0 +1,27 @@
+# RUN: %PYTHON %s | FileCheck %s
+# This is just a smoke test that the dialect is functional.
+from array import array
+
+from mlir.ir import *
+from mlir.dialects import ub
+from mlir.extras import types as T
+
+
+def constructAndPrintInModule(f):
+    print("\nTEST:", f.__name__)
+    with Context(), Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            f()
+        print(module)
+    return f
+
+
+# CHECK-LABEL: testSmoke
+@constructAndPrintInModule
+def testSmoke():
+    # CHECK: Value(%{{.*}} = ub.poison : f32
+    f32 = F32Type.get()
+    poison = ub.poison(f32)
+    print(poison)
+    assert isinstance(poison, Value)