[mlir][Python] Fix generation of accessors for Optional

[michalt]

Summary

Previously, in case there was only one Optional operand/result within
the list, we would always return None from the accessor, e.g., for a
single optional result we would generate:

return self.operation.results[0] if len(self.operation.results) > 1 else None

But what we really want is to return None only if the length of
results is smaller than the total number of element groups (i.e.,
the optional operand/result is in fact missing).

This commit also renames a few local variables in the generator to make
the distinction between isVariadic() and isVariableLength() a bit
more clear.

[michalt]

This is the first time I'm trying to do anything non-trivial with MLIR, so any feedback is very welcome (especially if I misunderstood/missed something :) Thanks!

[ftynse]

Thanks for the fix! Please let me know if you need help landing this.

mlir/test/python/dialects/python_test.py

132

Nit: spurious whitespace change (this file uses two newlines between functions AFAICS).

[michalt]

Fix whitespace

[michalt]

In D113855#3134009, @ftynse wrote:

Thanks for the fix! Please let me know if you need help landing this.

Thank you for a quick review! :) I don't have commit access, so it'd be great if you could land this, thanks!

[ftynse]

Something went wrong and now the diff contains only the whitespace changes. If you made them as two separate git commits, please indicate the two-commit range when updating the revision or squash them.

[michalt]

Correctly export the whatespace fix to Phab

[michalt]

In D113855#3137685, @ftynse wrote:

Something went wrong and now the diff contains only the whitespace changes. If you made them as two separate git commits, please indicate the two-commit range when updating the revision or squash them.

Yikes! I haven't used arc for a couple of years... and indeed I made a separate commit and didn't specify the range correctly when running arc diff. Sorry for that!

mlir/python/mlir/dialects/_linalg_ops_ext.py

op = self.build_generic(
         results=results,
         operands=[_get_op_result_or_value(o) for o in [value, output]],
         attributes=None,
         loc=loc,
         ip=ip)
     OpView.__init__(self, op)
     linalgDialect = Context.current.get_dialect_descriptor("linalg")
     fill_builtin_region(linalgDialect, self.operation)
-    # TODO: self.result is None. When len(results) == 1 we expect it to be
-    # results[0] as per _linalg_ops_gen.py. This seems like an orthogonal bug
-    # in the generator of _linalg_ops_gen.py where we have:
-    # ```
-    # def result(self):
-    #   return self.operation.results[0] \
-    #     if len(self.operation.results) > 1 else None
-    # ```
-
 
 class InitTensorOp:
   """Extends the linalg.init_tensor op."""
 
   def __init__(self,
                sizes: Union[Sequence[int], Sequence[Value]],
                element_type: Type,
                *,

mlir/test/mlir-tblgen/op-python-bindings.td

def OneOptionalOperandOp : TestOp<"one_optional_operand"> {
   // CHECK:     successors=_ods_successors, regions=regions, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def non_optional(self):
   // CHECK:   return self.operation.operands[0]
 
   // CHECK: @builtins.property
   // CHECK: def optional(self):
-  // CHECK:   return self.operation.operands[1] if len(self.operation.operands) > 2 else None
+  // CHECK:   return None if len(self.operation.operands) < 2 else self.operation.operands[1]
 
 }
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class OneVariadicOperandOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.one_variadic_operand"
 // CHECK-NOT: _ODS_OPERAND_SEGMENTS
 // CHECK-NOT: _ODS_RESULT_SEGMENTS

mlir/test/python/dialects/linalg/ops.py

with InsertionPoint(module.body):
       # CHECK-LABEL: func @fill_tensor
       #  CHECK-SAME:   %[[OUT:[0-9a-z]+]]: tensor<12x?xf32>
       #  CHECK-NEXT: %[[CST:.*]] = arith.constant 0.0{{.*}} : f32
       #  CHECK-NEXT: %[[RES:.*]] = linalg.fill(%[[CST]], %[[OUT]]) : f32, tensor<12x?xf32> -> tensor<12x?xf32>
       #  CHECK-NEXT: return %[[RES]] : tensor<12x?xf32>
       @builtin.FuncOp.from_py_func(RankedTensorType.get((12, -1), f32))
       def fill_tensor(out):
         zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.), result=f32).result
-        # TODO: FillOp.result is None. When len(results) == 1 we expect it to
-        # be results[0] as per _linalg_ops_gen.py. This seems like an
-        # orthogonal bug in the generator of _linalg_ops_gen.py.
-        return linalg.FillOp(output=out, value=zero).results[0]
+        return linalg.FillOp(output=out, value=zero).result
 
       # CHECK-LABEL: func @fill_buffer
       #  CHECK-SAME:   %[[OUT:[0-9a-z]+]]: memref<12x?xf32>
       #  CHECK-NEXT: %[[CST:.*]] = arith.constant 0.0{{.*}} : f32
       #  CHECK-NEXT: linalg.fill(%[[CST]], %[[OUT]]) : f32, memref<12x?xf32>
       #  CHECK-NEXT: return
       @builtin.FuncOp.from_py_func(MemRefType.get((12, -1), f32))
       def fill_buffer(out):

mlir/test/python/dialects/python_test.py

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: inferReturnTypes
 @run
 def inferReturnTypes():
   with Context() as ctx, Location.unknown(ctx):
     test.register_python_test_dialect(ctx)
     module = Module.create()
     with InsertionPoint(module.body):
       op = test.InferResultsOp()
with InsertionPoint(module.body):
 
       implied = test.InferResultsImpliedOp()
       # CHECK: i32
       print(implied.integer.type)
       # CHECK: f64
       print(implied.flt.type)
       # CHECK: index
       print(implied.index.type)
+
+
+# CHECK-LABEL: TEST: testOptionalOperandOp
+@run
+def testOptionalOperandOp():
+  with Context() as ctx, Location.unknown():
+    test.register_python_test_dialect(ctx)
+
+    module = Module.create()
+    with InsertionPoint(module.body):
+
+      op1 = test.OptionalOperandOp(None)
+      # CHECK: op1.input is None: True
+      print(f"op1.input is None: {op1.input is None}")
+
+      op2 = test.OptionalOperandOp(op1)
+      # CHECK: op2.input is None: False
+      print(f"op2.input is None: {op2.input is None}")

mlir/test/python/python_test_ops.td

 }
 
 def FirstAttrDeriveAttrOp : TestOp<"first_attr_derive_attr_op",
                                [FirstAttrDerivedResultType]> {
   let arguments = (ins AnyAttr:$iattr);
   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);
+}
+
 #endif // PYTHON_TEST_OPS

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

def {0}(self):
     return self.operation.{1}s[{3} + _ods_variadic_group_length - 1]
 )Py";
 
 /// Template for an optional element accessor:
 ///   {0} is the name of the accessor;
 ///   {1} is either 'operand' or 'result';
 ///   {2} is the total number of element groups;
 ///   {3} is the position of the current group in the group list.
+/// This works if we have only one variable-length group (and it's the optional
+/// operand/result): we can deduce it's absent if the `len(operation.{1}s)` is
+/// smaller than the total number of groups.
 constexpr const char *opOneOptionalTemplate = R"Py(
   @builtins.property
   def {0}(self):
-    return self.operation.{1}s[{3}] if len(self.operation.{1}s) > {2} else None
+    return None if len(self.operation.{1}s) < {2} else self.operation.{1}s[{3}]
 )Py";
 
 /// Template for the variadic group accessor in the single variadic group case:
 ///   {0} is the name of the accessor;
 ///   {1} is either 'operand' or 'result';
 ///   {2} is the total number of element groups;
 ///   {3} is the position of the current group in the group list.
 constexpr const char *opOneVariadicTemplate = R"Py(
return llvm::formatv("::mlir::OpTrait::AttrSized{0}{1}Segments",
                        llvm::StringRef(kind).drop_front());
 }
 
 /// Emits accessors to "elements" of an Op definition. Currently, the supported
 /// elements are operands and results, indicated by `kind`, which must be either
 /// `operand` or `result` and is used verbatim in the emitted code.
 static void emitElementAccessors(
     const Operator &op, raw_ostream &os, const char *kind,
-    llvm::function_ref<unsigned(const Operator &)> getNumVariadic,
+    llvm::function_ref<unsigned(const Operator &)> getNumVariableLength,
     llvm::function_ref<int(const Operator &)> getNumElements,
     llvm::function_ref<const NamedTypeConstraint &(const Operator &, int)>
         getElement) {
   assert(llvm::is_contained(
              llvm::SmallVector<StringRef, 2>{"operand", "result"}, kind) &&
          "unsupported kind");
 
   // Traits indicating how to process variadic elements.
   std::string sameSizeTrait =
       llvm::formatv("::mlir::OpTrait::SameVariadic{0}{1}Size",
                     llvm::StringRef(kind).take_front().upper(),
                     llvm::StringRef(kind).drop_front());
   std::string attrSizedTrait = attrSizedTraitForKind(kind);
 
-  unsigned numVariadic = getNumVariadic(op);
+  unsigned numVariableLength = getNumVariableLength(op);
 
-  // If there is only one variadic element group, its size can be inferred from
-  // the total number of elements. If there are none, the generation is
-  // straightforward.
-  if (numVariadic <= 1) {
+  // If there is only one variable-length element group, its size can be
+  // inferred from the total number of elements. If there are none, the
+  // generation is straightforward.
+  if (numVariableLength <= 1) {
     bool seenVariableLength = false;
     for (int i = 0, e = getNumElements(op); i < e; ++i) {
       const NamedTypeConstraint &element = getElement(op, i);
       if (element.isVariableLength())
         seenVariableLength = true;
       if (element.name.empty())
         continue;
       if (element.isVariableLength()) {
static void emitElementAccessors(
   // Handle the operations where variadic groups have the same size.
   if (op.getTrait(sameSizeTrait)) {
     int numPrecedingSimple = 0;
     int numPrecedingVariadic = 0;
     for (int i = 0, e = getNumElements(op); i < e; ++i) {
       const NamedTypeConstraint &element = getElement(op, i);
       if (!element.name.empty()) {
         os << llvm::formatv(opVariadicEqualPrefixTemplate,
-                            sanitizeName(element.name), kind, numVariadic,
+                            sanitizeName(element.name), kind, numVariableLength,
                             numPrecedingSimple, numPrecedingVariadic);
         os << llvm::formatv(element.isVariableLength()
                                 ? opVariadicEqualVariadicTemplate
                                 : opVariadicEqualSimpleTemplate,
                             kind);
       }
       if (element.isVariableLength())
         ++numPrecedingVariadic;

 }
 static int getNumResults(const Operator &op) { return op.getNumResults(); }
 static const NamedTypeConstraint &getResult(const Operator &op, int i) {
   return op.getResult(i);
 }
 
 /// Emits accessors to Op operands.
 static void emitOperandAccessors(const Operator &op, raw_ostream &os) {
-  auto getNumVariadic = [](const Operator &oper) {
+  auto getNumVariableLengthOperands = [](const Operator &oper) {
     return oper.getNumVariableLengthOperands();
   };
-  emitElementAccessors(op, os, "operand", getNumVariadic, getNumOperands,
-                       getOperand);
+  emitElementAccessors(op, os, "operand", getNumVariableLengthOperands,
+                       getNumOperands, getOperand);
 }
 
 /// Emits accessors Op results.
 static void emitResultAccessors(const Operator &op, raw_ostream &os) {
-  auto getNumVariadic = [](const Operator &oper) {
+  auto getNumVariableLengthResults = [](const Operator &oper) {
     return oper.getNumVariableLengthResults();
   };
-  emitElementAccessors(op, os, "result", getNumVariadic, getNumResults,
-                       getResult);
+  emitElementAccessors(op, os, "result", getNumVariableLengthResults,
+                       getNumResults, getResult);
 }
 
 /// Emits accessors to Op attributes.
 static void emitAttributeAccessors(const Operator &op,
                                    const AttributeClasses &attributeClasses,
                                    raw_ostream &os) {
   for (const auto &namedAttr : op.getAttributes()) {
     // Skip "derived" attributes because they are just C++ functions that we