[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

425

446

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
with Context() as ctx, Location.unknown(ctx):
         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)
+
+
+
+# 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)
 No newline at end of file

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,