[MLIR][python bindings] Fix inferReturnTypes + AttrSizedOperandSegments for optional operands

[makslevental]

Summary

Right now inferTypeOpInterface.inferReturnTypes fails because there's a cast in there to py::sequence which throws a TypeError when it tries to cast the Nones. Note Nones are inserted into operands for omitted operands passed to the generated builder:

operands.append(_get_op_result_or_value(start) if start is not None else None)
operands.append(_get_op_result_or_value(stop) if stop is not None else None)
operands.append(_get_op_result_or_value(step) if step is not None else None)

Note also that skipping appending to the list operands doesn't work either because build generic checks against the number of operand segments expected.

Currently the only way around is to handroll through ir.Operation.create.

[makslevental]

mlir/lib/Bindings/Python/IRInterfaces.cpp

72

This cast fails.

[rkayaith]

makes sense to me

[rkayaith]

mlir/test/python/dialects/python_test.py

383

441

can you fix the missing newline

[makslevental]

fix CHECK-LABEL

[makslevental]

format test

[makslevental]

format

mlir/lib/Bindings/Python/IRInterfaces.cpp

llvm::SmallVector<MlirValue> wrapOperands(std::optional<py::list> operandList) {
 
   if (!operandList || operandList->empty()) {
     return mlirOperands;
   }
 
   // Note: as the list may contain other lists this may not be final size.
   mlirOperands.reserve(operandList->size());
   for (const auto &&it : llvm::enumerate(*operandList)) {
+    if (it.value().is_none())
+      continue;
+
     PyValue *val;
     try {
       val = py::cast<PyValue *>(it.value());
       if (!val)
         throw py::cast_error();
       mlirOperands.push_back(val->get());
       continue;
     } catch (py::cast_error &err) {
       // Intentionally unhandled to try sequence below first.
       (void)err;
     }
 
     try {
       auto vals = py::cast<py::sequence>(it.value());
       for (py::object v : vals) {
         try {
           val = py::cast<PyValue *>(v);
           if (!val)
             throw py::cast_error();
           mlirOperands.push_back(val->get());
         } catch (py::cast_error &err) {
           throw py::value_error(

mlir/test/python/dialects/python_test.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 import mlir.dialects.func as func
 import mlir.dialects.python_test as test
 import mlir.dialects.tensor as tensor
+import mlir.dialects.arith as arith
+
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   return f
 
+
 # CHECK-LABEL: TEST: testAttributes
 @run
 def testAttributes():
   with Context() as ctx, Location.unknown():
     ctx.allow_unregistered_dialects = True
 
     #
     # Check op construction with attributes.
with Context() as ctx, Location.unknown():
     # CHECK: Unit: True
     op.unit = True
     print(f"Unit: {op.unit}")
 
     # CHECK: Unit: False
     del op.unit
     print(f"Unit: {op.unit}")
 
+
 # CHECK-LABEL: TEST: attrBuilder
 @run
 def attrBuilder():
   with Context() as ctx, Location.unknown():
     ctx.allow_unregistered_dialects = True
     op = test.AttributesOp(x_bool=True,
                            x_i16=1,
                            x_i32=2,
with InsertionPoint(module.body):
       print(same.two.type)
 
       first_type_attr = test.FirstAttrDeriveTypeAttrOp(
           inferred.results[1], TypeAttr.get(IndexType.get()))
       # CHECK-COUNT-2: index
       print(first_type_attr.one.type)
       print(first_type_attr.two.type)
 
-      first_attr = test.FirstAttrDeriveAttrOp(
-          FloatAttr.get(F32Type.get(), 3.14))
+      first_attr = test.FirstAttrDeriveAttrOp(FloatAttr.get(
+          F32Type.get(), 3.14))
       # CHECK-COUNT-3: f32
       print(first_attr.one.type)
       print(first_attr.two.type)
       print(first_attr.three.type)
 
       implied = test.InferResultsImpliedOp()
       # CHECK: i32
       print(implied.integer.type)

 # CHECK-LABEL: TEST: testTensorValue
 def testTensorValue():
   with Context() as ctx, Location.unknown():
     test.register_python_test_dialect(ctx)
 
     i8 = IntegerType.get_signless(8)
 
     class Tensor(test.TestTensorValue):
+
       def __str__(self):
         return super().__str__().replace("Value", "Tensor")
 
     module = Module.create()
     with InsertionPoint(module.body):
       t = tensor.EmptyOp([10, 10], i8).result
 
       # CHECK: Value(%{{.*}} = tensor.empty() : tensor<10x10xi8>)
with InsertionPoint(module.body):
       assert isinstance(test.TestTensorType.static_typeid, TypeID)
       # And it should be equal to the in-tree concrete type
       assert test.TestTensorType.static_typeid == t.type.typeid
 
 
 # CHECK-LABEL: TEST: inferReturnTypeComponents
 @run
 def inferReturnTypeComponents():
-    with Context() as ctx, Location.unknown(ctx):
-        test.register_python_test_dialect(ctx)
-        module = Module.create()
-        i32 = IntegerType.get_signless(32)
-        with InsertionPoint(module.body):
-            resultType = UnrankedTensorType.get(i32)
-            operandTypes = [
-                RankedTensorType.get([1, 3, 10, 10], i32),
-                UnrankedTensorType.get(i32),
-            ]
-            f = func.FuncOp(
-                "test_inferReturnTypeComponents", (operandTypes, [resultType])
-            )
-            entry_block = Block.create_at_start(f.operation.regions[0], operandTypes)
-            with InsertionPoint(entry_block):
-                ranked_op = test.InferShapedTypeComponentsOp(
-                    resultType, entry_block.arguments[0]
-                )
-                unranked_op = test.InferShapedTypeComponentsOp(
-                    resultType, entry_block.arguments[1]
-                )
+      f = func.FuncOp("test_inferReturnTypeComponents",
+                      (operandTypes, [resultType]))
+        ranked_op = test.InferShapedTypeComponentsOp(resultType,
+                                                     entry_block.arguments[0])
+        unranked_op = test.InferShapedTypeComponentsOp(resultType,
+                                                       entry_block.arguments[1])
 
-        # CHECK: has rank: True
-        # CHECK: rank: 4
-        # CHECK: element type: i32
-        # CHECK: shape: [1, 3, 10, 10]
-        iface = InferShapedTypeOpInterface(ranked_op)
-        shaped_type_components = iface.inferReturnTypeComponents(
-            operands=[ranked_op.operand]
-        )[0]
-        print("has rank:", shaped_type_components.has_rank)
-        print("rank:", shaped_type_components.rank)
-        print("element type:", shaped_type_components.element_type)
-        print("shape:", shaped_type_components.shape)
+        operands=[ranked_op.operand])[0]
 
-        # CHECK: has rank: False
-        # CHECK: rank: None
-        # CHECK: element type: i32
-        # CHECK: shape: None
-        iface = InferShapedTypeOpInterface(unranked_op)
-        shaped_type_components = iface.inferReturnTypeComponents(
-          operands=[unranked_op.operand]
-        )[0]
-        print("has rank:", shaped_type_components.has_rank)
-        print("rank:", shaped_type_components.rank)
-        print("element type:", shaped_type_components.element_type)
-        print("shape:", shaped_type_components.shape)
+        operands=[unranked_op.operand])[0]
+
+
+# CHECK-LABEL: TEST: testInferTypeOpInterface
+@run
+def testInferTypeOpInterface():
+  with Context() as ctx, Location.unknown(ctx):
+    test.register_python_test_dialect(ctx)
+    module = Module.create()
+    with InsertionPoint(module.body):
+      i64 = IntegerType.get_signless(64)
+      zero = arith.ConstantOp(i64, 0)
+
+      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)

mlir/test/python/python_test_ops.td

static ::mlir::LogicalResult inferReturnTypes(
       ::mlir::Builder b(context);
       inferredReturnTypes.push_back(b.getI32Type());
       inferredReturnTypes.push_back(b.getI64Type());
       return ::mlir::success();
     }
   }];
 }
 
+def I32OrF32 : TypeConstraint<Or<[I32.predicate, F32.predicate]>,
+                                 "i32 or f32">;
+
+def InferResultsVariadicInputsOp : TestOp<"infer_results_variadic_inputs_op",
+    [InferTypeOpInterface, AttrSizedOperandSegments]> {
+  let arguments = (ins Optional<I64>:$single, Optional<I64>:$doubled);
+  let results = (outs I32OrF32:$res);
+
+  let extraClassDeclaration = [{
+    static ::mlir::LogicalResult inferReturnTypes(
+      ::mlir::MLIRContext *context, ::std::optional<::mlir::Location> location,
+      ::mlir::ValueRange operands, ::mlir::DictionaryAttr attributes,
+      ::mlir::OpaqueProperties,
+      ::mlir::RegionRange regions,
+      ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
+      ::mlir::Builder b(context);
+      if (operands.size() == 1)
+          inferredReturnTypes.push_back(b.getI32Type());
+      else if (operands.size() == 2)
+          inferredReturnTypes.push_back(b.getF32Type());
+      return ::mlir::success();
+    }
+  }];
+}
+
 // If all result types are buildable, the InferTypeOpInterface is implied and is
 // autogenerated by C++ ODS.
 def InferResultsImpliedOp : TestOp<"infer_results_implied_op"> {
   let results = (outs I32:$integer, F64:$flt, Index:$index);
 }
 
 def InferShapedTypeComponentsOp : TestOp<"infer_shaped_type_components_op",
   [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,