[mlir][py] Fix infer return type invocation for variadics

[jpienaar]

Summary

Previously we only allowed the flattened list passed in, but the same
input provided here as to buildGeneric so flatten accordingly. We have
less info here than in buildGeneric so the error is more generic if
unpacking fails.

[stella.stamenova]

Looks like this caused a failure in the mlir buildbot due to an unused variable: https://lab.llvm.org/buildbot/#/builders/13/builds/31682

[jpienaar]

Sending quick fix.

• smime.p7s3 KBDownload

mlir/lib/Bindings/Python/IRInterfaces.cpp

 //===- IRInterfaces.cpp - MLIR IR interfaces pybind -----------------------===//
 //
 // 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 <utility>
 #include <optional>
 
 #include "IRModule.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/Interfaces.h"
+#include "llvm/ADT/STLExtras.h"
 
 namespace py = pybind11;
 
 namespace mlir {
 namespace python {
 
 constexpr static const char *constructorDoc =
     R"(Creates an interface from a given operation/opview object or from a
static void appendResultsCallback(intptr_t nTypes, MlirType *types,
     auto *data = static_cast<AppendResultsCallbackData *>(userData);
     data->inferredTypes.reserve(data->inferredTypes.size() + nTypes);
     for (intptr_t i = 0; i < nTypes; ++i) {
       data->inferredTypes.emplace_back(data->pyMlirContext.getRef(), types[i]);
     }
   }
 
   /// Given the arguments required to build an operation, attempts to infer its
-  /// return types. Throws value_error on faliure.
+  /// return types. Throws value_error on failure.
   std::vector<PyType>
-  inferReturnTypes(std::optional<std::vector<PyValue>> operands,
+  inferReturnTypes(std::optional<py::list> operandList,
                    std::optional<PyAttribute> attributes,
                    std::optional<std::vector<PyRegion>> regions,
                    DefaultingPyMlirContext context,
                    DefaultingPyLocation location) {
     llvm::SmallVector<MlirValue> mlirOperands;
     llvm::SmallVector<MlirRegion> mlirRegions;
 
-    if (operands) {
-      mlirOperands.reserve(operands->size());
-      for (PyValue &value : *operands) {
-        mlirOperands.push_back(value);
+    if (operandList && !operandList->empty()) {
+      // 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)) {
+        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) {
+        }
+
+        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(
+                  (llvm::Twine("Operand ") + llvm::Twine(it.index()) +
+                   " must be a Value or Sequence of Values (" + err.what() +
+                   ")")
+                      .str());
+            }
+          }
+          continue;
+        } catch (py::cast_error &err) {
+          throw py::value_error(
+              (llvm::Twine("Operand ") + llvm::Twine(it.index()) +
+               " must be a Value or Sequence of Values (" + err.what() + ")")
+                  .str());
+        }
+
+        throw py::cast_error();
       }
     }
 
     if (regions) {
       mlirRegions.reserve(regions->size());
       for (PyRegion &region : *regions) {
         mlirRegions.push_back(region);
       }

mlir/lib/Interfaces/ViewLikeInterface.cpp

 
 LogicalResult mlir::verifyListOfOperandsOrIntegers(Operation *op,
                                                    StringRef name,
                                                    unsigned numElements,
                                                    ArrayRef<int64_t> staticVals,
                                                    ValueRange values) {
   // Check static and dynamic offsets/sizes/strides does not overflow type.
   if (staticVals.size() != numElements)
-    return op->emitError("expected ")
-           << numElements << " " << name << " values";
+    return op->emitError("expected ") << numElements << " " << name
+                                      << " values, got " << staticVals.size();
   unsigned expectedNumDynamicEntries =
       llvm::count_if(staticVals, [&](int64_t staticVal) {
         return ShapedType::isDynamic(staticVal);
       });
   if (values.size() != expectedNumDynamicEntries)
     return op->emitError("expected ")
            << expectedNumDynamicEntries << " dynamic " << name << " values";
   return success();

mlir/python/mlir/_mlir_libs/_mlir/ir.pyi

class IndexType(Type):
     def __init__(self, cast_from_type: Type) -> None: ...
     @staticmethod
     def get(*args, **kwargs) -> IndexType: ...
     @staticmethod
     def isinstance(arg: Any) -> bool: ...
 
 class InferTypeOpInterface:
     def __init__(self, object: object, context: Optional[Context] = None) -> None: ...
-    def inferReturnTypes(self, operands: Optional[List[Value]] = None, attributes: Optional[Attribute] = None, regions: Optional[List[Region]] = None, context: Optional[Context] = None, loc: Optional[Location] = None) -> List[Type]: ...
+    def inferReturnTypes(self, operands: Optional[List] = None, attributes: Optional[Attribute] = None, regions: Optional[List[Region]] = None, context: Optional[Context] = None, loc: Optional[Location] = None) -> List[Type]: ...
     @property
     def operation(self) -> Operation: ...
     @property
     def opview(self) -> OpView: ...
 
 class InsertionPoint:
     current: ClassVar[InsertionPoint] = ...  # read-only
     @overload

mlir/test/python/dialects/tensor.py

with InsertionPoint(module.body):
 
       # CHECK-LABEL: func @zero_d
       # CHECK: %0 = tensor.empty() : tensor<f32>
       @func.FuncOp.from_py_func()
       def zero_d():
         return tensor.EmptyOp([], f32)
 
   print(module)
+
+
+# CHECK-LABEL: TEST: testInferTypesInsertSlice
+@run
+def testInferTypesInsertSlice():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    f32Type = F32Type.get()
+    indexType = IndexType.get()
+    with InsertionPoint(module.body):
+
+      @func.FuncOp.from_py_func(
+          RankedTensorType.get((1, 1), f32Type),
+          RankedTensorType.get((1, 1), f32Type))
+      # CHECK: func @f
+      # CHECK:      tensor.insert_slice %arg0 into %arg1[0, 0] [1, 1] [0, 0] :
+      # CHECK-SAME:   tensor<1x1xf32> into tensor<1x1xf32>
+      def f(source, dest):
+        c0 = arith.ConstantOp(indexType, 0)
+        c1 = arith.ConstantOp(indexType, 1)
+        d0 = tensor.InsertSliceOp(source, dest, [], [], [],
+                                  DenseI64ArrayAttr.get([0, 0]),
+                                  DenseI64ArrayAttr.get([1, 1]),
+                                  DenseI64ArrayAttr.get([0, 0]))
+        return [d0.result]
+
+  print(module)