[MLIR][python bindings] I32BitEnumAttrs

[makslevental]

Summary

Depends on D157934

[makslevental]

mlir/test/python/lib/PythonTestDialect.h

28–47

these two things need to be in tablegen rather than by hand

[makslevental]

check IntFlag

mlir/python/CMakeLists.txt

declare_mlir_python_sources(
     ADD_TO_PARENT MLIRPythonTestSources.Dialects
     SOURCES dialects/python_test.py)
   set(LLVM_TARGET_DEFINITIONS
     "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
   mlir_tablegen(
     "dialects/_python_test_ops_gen.py"
     -gen-python-op-bindings
     -bind-dialect=python_test)
+  mlir_tablegen(
+    "dialects/_python_test_enum_gen.py"
+    -gen-python-enum-bindings)
   add_public_tablegen_target(PythonTestDialectPyIncGen)
   declare_mlir_python_sources(
     MLIRPythonTestSources.Dialects.PythonTest.ops_gen
     ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
     ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
-    SOURCES "dialects/_python_test_ops_gen.py")
+    SOURCES
+      "dialects/_python_test_ops_gen.py"
+      "dialects/_python_test_enum_gen.py"
+    )
 
   declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
     MODULE_NAME _mlirPythonTest
     ADD_TO_PARENT MLIRPythonTestSources.Dialects
     ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
     SOURCES
       PythonTestModule.cpp
     PRIVATE_LINK_LIBS

mlir/python/mlir/dialects/python_test.py

 #  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 ._python_test_ops_gen import *
+from ._python_test_enum_gen import *
 from .._mlir_libs._mlirPythonTest import TestAttr, TestType, TestTensorValue, TestIntegerRankedTensorType
 
 
 def register_python_test_dialect(context, load=True):
     from .._mlir_libs import _mlirPythonTest
 
     _mlirPythonTest.register_python_test_dialect(context, load)

mlir/test/python/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS python_test_ops.td)
 mlir_tablegen(lib/PythonTestDialect.h.inc -gen-dialect-decls)
 mlir_tablegen(lib/PythonTestDialect.cpp.inc -gen-dialect-defs)
 mlir_tablegen(lib/PythonTestOps.h.inc -gen-op-decls)
 mlir_tablegen(lib/PythonTestOps.cpp.inc -gen-op-defs)
 mlir_tablegen(lib/PythonTestAttributes.h.inc -gen-attrdef-decls)
 mlir_tablegen(lib/PythonTestAttributes.cpp.inc -gen-attrdef-defs)
 mlir_tablegen(lib/PythonTestTypes.h.inc -gen-typedef-decls)
 mlir_tablegen(lib/PythonTestTypes.cpp.inc -gen-typedef-defs)
+mlir_tablegen(lib/PythonTestEnums.h.inc -gen-enum-decls)
+mlir_tablegen(lib/PythonTestEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRPythonTestIncGen)
 
 add_subdirectory(lib)

mlir/test/python/dialects/python_test.py

with Context() as ctx, Location.unknown(ctx):
 
             one_operand = test.InferResultsVariadicInputsOp(single=zero, doubled=None)
             # CHECK: i32
             print(one_operand.result.type)
 
             two_operands = test.InferResultsVariadicInputsOp(single=zero, doubled=zero)
             # CHECK: f32
             print(two_operands.result.type)
+
+
+# CHECK-LABEL: TEST: testBitEnum
+@run
+def testBitEnum():
+    with Context() as ctx, Location.unknown(ctx):
+        test.register_python_test_dialect(ctx)
+        module = Module.create()
+        with InsertionPoint(module.body):
+            op = test.OpWithBareBitEnumVerticalBar(1)
+            # CHECK: python_test.op_with_bit_bar_enum_vbar user
+            print(op)
+            op = test.OpWithBareBitEnumVerticalBar(2)
+            # CHECK: python_test.op_with_bit_bar_enum_vbar group
+            print(op)
+            op = test.OpWithBareBitEnumVerticalBar(1 | 2)
+            # CHECK: python_test.op_with_bit_bar_enum_vbar user | group
+            print(op)
+        print(module.body)

mlir/test/python/lib/PythonTestDialect.h

 //===- PythonTestDialect.h - PythonTest dialect definition ------*- C++ -*-===//
 //
 // 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 MLIR_TEST_PYTHON_LIB_PYTHONTESTDIALECT_H
 #define MLIR_TEST_PYTHON_LIB_PYTHONTESTDIALECT_H
 
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "llvm/ADT/StringExtras.h"
 
 #include "PythonTestDialect.h.inc"
 
 #define GET_OP_CLASSES
 #include "PythonTestOps.h.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "PythonTestAttributes.h.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "PythonTestTypes.h.inc"
 
+namespace python_test {
+inline std::string stringifyTestBitEnumVerticalBar(llvm::APInt symbol) {
+  auto val = static_cast<uint32_t>(symbol.getLimitedValue());
+  assert(7u == (7u | val) && "invalid bits set in bit enum");
+  ::llvm::SmallVector<::llvm::StringRef, 2> strs;
+
+  if (1u == (1u & val))
+    strs.push_back("user");
+
+  if (2u == (2u & val))
+    strs.push_back("group");
+
+  if (4u == (4u & val))
+    strs.push_back("other");
+  return llvm::join(strs, " | ");
+}
+// ::mlir::ParseResult OpWithBareBitEnumVerticalBar::parse ...
+// valueAttr =
+// parser.getBuilder().getIntegerAttr(parser.getBuilder().getIntegerType(32),
+// static_cast<uint32_t>(*attrOptional));
+
+} // namespace python_test
+
+#include "PythonTestEnums.h.inc"
+
 #endif // MLIR_TEST_PYTHON_LIB_PYTHONTESTDIALECT_H

mlir/test/python/lib/PythonTestDialect.cpp

 //===- PythonTestDialect.cpp - PythonTest dialect definition --------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "PythonTestDialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 #include "PythonTestDialect.cpp.inc"
+#include "PythonTestEnums.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "PythonTestAttributes.cpp.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "PythonTestTypes.cpp.inc"
 
 #define GET_OP_CLASSES

mlir/test/python/python_test_ops.td

 //===-- python_test_ops.td - Python test Op definitions ----*- 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_TEST_OPS
 #define PYTHON_TEST_OPS
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/OpBase.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/IR/EnumAttr.td"
 
 def Python_Test_Dialect : Dialect {
   let name = "python_test";
   let cppNamespace = "python_test";
 
   let useDefaultTypePrinterParser = 1;
   let useDefaultAttributePrinterParser = 1;
 }
def FirstAttrDeriveAttrOp : TestOp<"first_attr_derive_attr_op",
   let results = (outs AnyType:$one, AnyType:$two, AnyType:$three);
 }
 
 def OptionalOperandOp : TestOp<"optional_operand_op"> {
   let arguments = (ins Optional<AnyType>:$input);
   let results = (outs I32:$result);
 }
 
+// Define an enum with a different separator
+def TestBitEnumVerticalBar
+    : I32BitEnumAttr<"TestBitEnumVerticalBar", "another test bit enum", [
+        I32BitEnumAttrCaseBit<"User", 0, "user">,
+        I32BitEnumAttrCaseBit<"Group", 1, "group">,
+        I32BitEnumAttrCaseBit<"Other", 2, "other">,
+      ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "python_test";
+  let separator = " | ";
+}
+
+// Define an op that contains the bare enum attribute.
+def OpWithBareBitEnumVerticalBar : TestOp<"op_with_bit_bar_enum_vbar"> {
+  let arguments = (ins TestBitEnumVerticalBar:$value,
+                   OptionalAttr<AnyAttr>:$tag);
+  let assemblyFormat = "$value (`tag` $tag^)? attr-dict";
+}
+
 #endif // PYTHON_TEST_OPS