inferReturnTypes triggers a failure in tensor.InsertSliceOp

muneeb · January 27, 2023, 4:37pm

A recent change upstream under the following commit introduces the logic to infer the result type based on a number of parameters such as those used by tensor.InsertSliceOp

[mlir][ods] Generate inferReturnTypes for ops with TypesMatchWith

1b60f0d73c34fec4648bb05f98db75008a50f4d8

The use of infer

The following code reproduces the bug

from mlir.ir import *
import mlir.dialects.func as func
import mlir.dialects.tensor as tensor
 
 
def run(f):
  f()
  return f
 
 
@run
def testInsertSliceOp():
  with Context() as ctx, Location.unknown():
    module = Module.create()
    f32Type = F32Type.get()
    source_ty = RankedTensorType.get([1, 1], f32Type)
    dest_ty = RankedTensorType.get([8, 8], f32Type)
    with InsertionPoint(module.body):
      # CHECK-LABEL: func @insert_slice
      # CHECK: (%[[SRC:.+]]: tensor<1x1xf32>, %[[DST:.+]]: tensor<8x8xf32>) -> tensor<8x8xf32> {
      # CHECK: %[[T:.+]] = tensor.insert_slice %[[SRC]] into %[[DST]][0, 0] [1, 1] [1, 1] : tensor<1x1xf32> into tensor<8x8xf32>
      # CHECK return %[[T]] : tensor<8x8xf32>
      @func.FuncOp.from_py_func(source_ty, dest_ty)
      def insert_slice(src, dst):
        insert_slice = tensor.InsertSliceOp(
            source=src,
            dest=dst,
            offsets=[],
            sizes=[],
            strides=[],
            static_offsets=DenseI64ArrayAttr.get([0, 0]),
            static_sizes=DenseI64ArrayAttr.get([1, 1]),
            static_strides=DenseI64ArrayAttr.get([1, 1]))
        return [insert_slice.result]
    print(module)

stevenvar · February 1, 2023, 4:06pm

@Mogball , it seems you’re the person who worked on the commit that added type inference for operations with the TypesMatchWith trait. I hope you don’t mind us checking with you for confirmation that this is an actual issue?

For a quick introduction, @muneeb and I are both compiler engineers in Huawei Cambridge (UK) working with MLIR, and we noticed this issue recently after a rebase of our internal fork of LLVM against upstream.

The problem seems to be related to the fact that some operands of InsertSliceOp (offsets, sizes and strides) are lists of Value, while the type inference method (inferReturnTypes) actually expects each operand to be just a single Value, hence failing with a type error when called:

TypeError: inferReturnTypes(): incompatible function arguments. The following argument types are supported:
    1. (self: mlir._mlir_libs._mlir.ir.InferTypeOpInterface, operands: Optional[List[mlir._mlir_libs._mlir.ir.Value]] = None, attributes: Optional[mlir._mlir_libs._mlir.ir.Attribute] = None, regions: Optional[List[mlir._mlir_libs._mlir.ir.Region]] = None, context: mlir.ir.Context = None, loc: mlir.ir.Location = None) -> List[mlir._mlir_libs._mlir.ir.Type]

jpienaar · February 1, 2023, 4:15pm

The python bindings stops requiring the type if the result type can be inferred which means as we refine type inference or improve modelling python code breaks (we have a discussion about this and alternative proposed in ⚙ D126238 [mlir][python] Move results types to keyword-only arg. ). In this case you should just omit the result type.

stevenvar · February 1, 2023, 4:19pm

The issue is not that we’re passing too many parameters to the Operation, it is that even by omitting the result type (as it is already done in the example posted in the original post) the call to the method that does the type inference (introduced by [mlir][ods] Generate inferReturnTypes for ops with TypesMatchWith ) fails

jpienaar · February 1, 2023, 4:36pm

I asked about variadics and told “If the op has variadic results, type inference is not generated”, seems we need to restrict on operands too for typesmatchwith (I think we need to change the inference method to use adaptors to be able to handle variadics better).

stevenvar · February 1, 2023, 4:44pm

Interesting. Our in-house hacky fix was to add code that filters out every operand that is not of type Value before calling the type inference function (since these operands are not involved in the inference of the output type anyway) but it felt it was far from an acceptable long-term solution.

jpienaar · February 1, 2023, 4:48pm

Ouch, ok let’s see if we can fix this And obviously we need more tests here upstream. I’m OOO but will try to restrict this for now/make it more conservative and then we can expand support with tests.

stevenvar · February 1, 2023, 4:50pm

Sounds great, thanks a lot

jpienaar · February 2, 2023, 5:25am

This seems to be a different issue:

github.com

llvm/llvm-project/blob/4f2fde781fa93294c480236392826dd1d2cb0741/mlir/lib/Bindings/Python/IRInterfaces.cpp#L188


      
            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.
          std::vector<PyType>
          inferReturnTypes(std::optional<std::vector<PyValue>> operands,
                           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) {

defines the input to be a vector of values, but for variadics one can have lists of lists, which is why in

github.com

llvm/llvm-project/blob/4f2fde781fa93294c480236392826dd1d2cb0741/mlir/lib/Bindings/Python/IRCore.cpp#L1363


      
            mlirOperationDestroy(operation);
            valid = false;
          }
          
          
//------------------------------------------------------------------------------
          // PyOpView
          //------------------------------------------------------------------------------
          
          
py::object
          PyOpView::buildGeneric(const py::object &cls, py::list resultTypeList,
                                 py::list operandList, std::optional<py::dict> attributes,
                                 std::optional<std::vector<PyBlock *>> successors,
                                 std::optional<int> regions,
                                 DefaultingPyLocation location,
                                 const py::object &maybeIp) {
            PyMlirContextRef context = location->getContext();
            // Class level operation construction metadata.
            std::string name = py::cast<std::string>(cls.attr("OPERATION_NAME"));
            // Operand and result segment specs are either none, which does no
            // variadic unpacking, or a list of ints with segment sizes, where each
            // element is either a positive number (typically 1 for a scalar) or -1 to

one has the unpacking, and note in the generated code both the inferReturnTypes and build_generic take the same operands as input.

The change made just exposed this but was not the cause. Do you mind filing a github issue for this? (I did not check if my original proposal would work around this, but it would only be a workaround).

jpienaar · February 7, 2023, 3:07am

Submitted change that should enable this, let me know if it works, else we should perhaps follow up via github issue.

stevenvar · February 8, 2023, 4:16pm

Sorry for the delay, yes it works fine now, that’s great

Many thanks!