[mlir][func] Add support for nested tuples to TestDecomposeCallGraphTypes.

[ingomueller-net]

Summary

Nested tuples were only supported in some narrow edge cases (and
potentially only because the test ops like test.make_tuple aren't
properly verified). This patch adds a couple of test cases with tested
tuple types and makes them work in the test pass by extending the
argument materialization and decomposition functions.

[ingomueller-net]

Minor updates: remove some dead code from test, explain test cases better, replace term "field" with element in comments and variable names.

[ingomueller-net]

mlir/include/mlir/Dialect/Func/Transforms/DecomposeCallGraphTypes.h

57

I added std::decay like it is done for TypeConverter::addConversion in order to allow passing function pointers in addition to lambdas.

mlir/include/mlir/Dialect/Func/Transforms/DecomposeCallGraphTypes.h

public:
   /// This method tries to decompose a value of a certain type using provided
   /// decompose callback functions. If it is unable to do so, the original value
   /// is returned.
   void decomposeValue(OpBuilder &, Location, Type, Value,
                       SmallVectorImpl<Value> &);
 
   /// This method registers a callback function that will be called to decompose
   /// a value of a certain type into 0, 1, or multiple values.
-  template <typename FnT,
-            typename T = typename llvm::function_traits<FnT>::template arg_t<2>>
+  template <typename FnT, typename T = typename llvm::function_traits<
+                              std::decay_t<FnT>>::template arg_t<2>>
   void addDecomposeValueConversion(FnT &&callback) {
     decomposeValueConversions.emplace_back(
         wrapDecomposeValueConversionCallback<T>(std::forward<FnT>(callback)));
   }
 
 private:
   using DecomposeValueConversionCallFn =
       std::function<std::optional<LogicalResult>(

mlir/test/Transforms/decompose-call-graph-types.mlir

 // CHECK-SAME:                                   %[[ARG0:.*]]: i1) -> i1 {
 // CHECK:           return %[[ARG0]] : i1
 func.func @recursive_decomposition(%arg0: tuple<tuple<tuple<i1>>>) -> tuple<tuple<tuple<i1>>> {
   return %arg0 : tuple<tuple<tuple<i1>>>
 }
 
 // -----
 
+// Test case: Type that needs to be recursively decomposed at different recursion depths.
+
+// CHECK-LABEL:   func @mixed_recursive_decomposition(
+// CHECK-SAME:                 %[[ARG0:.*]]: i1,
+// CHECK-SAME:                 %[[ARG1:.*]]: i2) -> (i1, i2) {
+// CHECK:           %[[V0:.*]] = "test.make_tuple"() : () -> tuple<>
+// CHECK:           %[[V1:.*]] = "test.make_tuple"(%[[ARG0]]) : (i1) -> tuple<i1>
+// CHECK:           %[[V2:.*]] = "test.make_tuple"(%[[ARG1]]) : (i2) -> tuple<i2>
+// CHECK:           %[[V3:.*]] = "test.make_tuple"(%[[V2]]) : (tuple<i2>) -> tuple<tuple<i2>>
+// CHECK:           %[[V4:.*]] = "test.make_tuple"(%[[V0]], %[[V1]], %[[V3]]) : (tuple<>, tuple<i1>, tuple<tuple<i2>>) -> tuple<tuple<>, tuple<i1>, tuple<tuple<i2>>>
+// CHECK:           %[[V5:.*]] = "test.get_tuple_element"(%[[V4]]) {index = 0 : i32} : (tuple<tuple<>, tuple<i1>, tuple<tuple<i2>>>) -> tuple<>
+// CHECK:           %[[V6:.*]] = "test.get_tuple_element"(%[[V4]]) {index = 1 : i32} : (tuple<tuple<>, tuple<i1>, tuple<tuple<i2>>>) -> tuple<i1>
+// CHECK:           %[[V7:.*]] = "test.get_tuple_element"(%[[V6]]) {index = 0 : i32} : (tuple<i1>) -> i1
+// CHECK:           %[[V8:.*]] = "test.get_tuple_element"(%[[V4]]) {index = 2 : i32} : (tuple<tuple<>, tuple<i1>, tuple<tuple<i2>>>) -> tuple<tuple<i2>>
+// CHECK:           %[[V9:.*]] = "test.get_tuple_element"(%[[V8]]) {index = 0 : i32} : (tuple<tuple<i2>>) -> tuple<i2>
+// CHECK:           %[[V10:.*]] = "test.get_tuple_element"(%[[V9]]) {index = 0 : i32} : (tuple<i2>) -> i2
+// CHECK:           return %[[V7]], %[[V10]] : i1, i2
+func.func @mixed_recursive_decomposition(%arg0: tuple<tuple<>, tuple<i1>, tuple<tuple<i2>>>) -> tuple<tuple<>, tuple<i1>, tuple<tuple<i2>>> {
+  return %arg0 : tuple<tuple<>, tuple<i1>, tuple<tuple<i2>>>
+}
+
+// -----
+
 // Test case: Check decomposition of calls.
 
 // CHECK-LABEL:   func private @callee(i1, i32) -> (i1, i32)
 func.func private @callee(tuple<i1, i32>) -> tuple<i1, i32>
 
 // CHECK-LABEL:   func @caller(
 // CHECK-SAME:                 %[[ARG0:.*]]: i1,
 // CHECK-SAME:                 %[[ARG1:.*]]: i32) -> (i1, i32) {

 // CHECK:           return %[[RET0]], %[[RET1]] : i1, i32
 func.func @unconverted_op_result() -> tuple<i1, i32> {
   %0 = "test.source"() : () -> (tuple<i1, i32>)
   return %0 : tuple<i1, i32>
 }
 
 // -----
 
+// Test case: Ensure decompositions are inserted properly around results of
+// unconverted ops in the case of different nesting levels.
+
+// CHECK-LABEL:   func @nested_unconverted_op_result(
+// CHECK-SAME:                 %[[ARG0:.*]]: i1,
+// CHECK-SAME:                 %[[ARG1:.*]]: i32) -> (i1, i32) {
+// CHECK:           %[[V0:.*]] = "test.make_tuple"(%[[ARG1]]) : (i32) -> tuple<i32>
+// CHECK:           %[[V1:.*]] = "test.make_tuple"(%[[ARG0]], %[[V0]]) : (i1, tuple<i32>) -> tuple<i1, tuple<i32>>
+// CHECK:           %[[V2:.*]] = "test.op"(%[[V1]]) : (tuple<i1, tuple<i32>>) -> tuple<i1, tuple<i32>>
+// CHECK:           %[[V3:.*]] = "test.get_tuple_element"(%[[V2]]) {index = 0 : i32} : (tuple<i1, tuple<i32>>) -> i1
+// CHECK:           %[[V4:.*]] = "test.get_tuple_element"(%[[V2]]) {index = 1 : i32} : (tuple<i1, tuple<i32>>) -> tuple<i32>
+// CHECK:           %[[V5:.*]] = "test.get_tuple_element"(%[[V4]]) {index = 0 : i32} : (tuple<i32>) -> i32
+// CHECK:           return %[[V3]], %[[V5]] : i1, i32
+func.func @nested_unconverted_op_result(%arg: tuple<i1, tuple<i32>>) -> tuple<i1, tuple<i32>> {
+  %0 = "test.op"(%arg) : (tuple<i1, tuple<i32>>) -> (tuple<i1, tuple<i32>>)
+  return %0 : tuple<i1, tuple<i32>>
+}
+
+// -----
+
 // Test case: Check mixed decomposed and non-decomposed args.
 // This makes sure to test the cases if 1:0, 1:1, and 1:N decompositions.
 
 // CHECK-LABEL:   func private @callee(i1, i2, i3, i4, i5, i6) -> (i1, i2, i3, i4, i5, i6)
 func.func private @callee(tuple<>, i1, tuple<i2>, i3, tuple<i4, i5>, i6) -> (tuple<>, i1, tuple<i2>, i3, tuple<i4, i5>, i6)
 
 // CHECK-LABEL:   func @caller(
 // CHECK-SAME:                 %[[I1:.*]]: i1,

mlir/test/lib/Dialect/Func/TestDecomposeCallGraphTypes.cpp

 #include "mlir/Dialect/Func/Transforms/DecomposeCallGraphTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 namespace {
+/// Creates a sequence of `test.get_tuple_element` ops for all elements of a
+/// given tuple value. If some tuple elements are, in turn, tuples, the elements
+/// of those are extracted recursively such that the returned values have the
+/// same types as `resultTypes.getFlattenedTypes()`.
+static LogicalResult buildDecomposeTuple(OpBuilder &builder, Location loc,
+                                         TupleType resultType, Value value,
+                                         SmallVectorImpl<Value> &values) {
+  for (unsigned i = 0, e = resultType.size(); i < e; ++i) {
+    Type elementType = resultType.getType(i);
+    Value element = builder.create<test::GetTupleElementOp>(
+        loc, elementType, value, builder.getI32IntegerAttr(i));
+    if (auto nestedTupleType = elementType.dyn_cast<TupleType>()) {
+      // Recurse if the current element is also a tuple.
+      if (failed(buildDecomposeTuple(builder, loc, nestedTupleType, element,
+                                     values)))
+        return failure();
+    } else {
+      values.push_back(element);
+    }
+  }
+  return success();
+}
+
+/// Creates a `test.make_tuple` op out of the given inputs building a tuple of
+/// type `resultType`. If that type is nested, each nested tuple is built
+/// recursively with another `test.make_tuple` op.
+static std::optional<Value> buildMakeTupleOp(OpBuilder &builder,
+                                             TupleType resultType,
+                                             ValueRange inputs, Location loc) {
+  // Build one value for each element at this nesting level.
+  SmallVector<Value> elements;
+  elements.reserve(resultType.getTypes().size());
+  ValueRange::iterator inputIt = inputs.begin();
+  for (Type elementType : resultType.getTypes()) {
+    if (auto nestedTupleType = elementType.dyn_cast<TupleType>()) {
+      // Determine how many input values are needed for the nested elements of
+      // the nested TupleType and advance inputIt by that number.
+      // TODO: We only need the *number* of nested types, not the types itself.
+      //       Maybe it's worth adding a more efficient overload?
+      SmallVector<Type> nestedFlattenedTypes;
+      nestedTupleType.getFlattenedTypes(nestedFlattenedTypes);
+      size_t numNestedFlattenedTypes = nestedFlattenedTypes.size();
+      ValueRange nestedFlattenedelements(inputIt,
+                                         inputIt + numNestedFlattenedTypes);
+      inputIt += numNestedFlattenedTypes;
+
+      // Recurse on the values for the nested TupleType.
+      std::optional<Value> res = buildMakeTupleOp(builder, nestedTupleType,
+                                                  nestedFlattenedelements, loc);
+      if (!res.has_value())
+        return {};
+
+      // The tuple constructed by the conversion is the element value.
+      elements.push_back(res.value());
+    } else {
+      // Base case: take one input as is.
+      elements.push_back(*inputIt++);
+    }
+  }
+
+  // Assemble the tuple from the elements.
+  return builder.create<test::MakeTupleOp>(loc, resultType, elements);
+}
+
 /// A pass for testing call graph type decomposition.
 ///
 /// This instantiates the patterns with a TypeConverter and ValueDecomposer
 /// that splits tuple types into their respective element types.
 /// For example, `tuple<T1, T2, T3> --> T1, T2, T3`.
 struct TestDecomposeCallGraphTypes
     : public PassWrapper<TestDecomposeCallGraphTypes, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDecomposeCallGraphTypes)
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<test::TestDialect>();
   }
   StringRef getArgument() const final {
     return "test-decompose-call-graph-types";
   }
   StringRef getDescription() const final {
     return "Decomposes types at call graph boundaries.";
   }
   void runOnOperation() override {
     ModuleOp module = getOperation();
     auto *context = &getContext();
     TypeConverter typeConverter;
     ConversionTarget target(*context);
-    ValueDecomposer decomposer;
     RewritePatternSet patterns(context);
 
     target.addLegalDialect<test::TestDialect>();
 
     target.addDynamicallyLegalOp<func::ReturnOp>([&](func::ReturnOp op) {
       return typeConverter.isLegal(op.getOperandTypes());
     });
     target.addDynamicallyLegalOp<func::CallOp>(
         [&](func::CallOp op) { return typeConverter.isLegal(op); });
     target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
       return typeConverter.isSignatureLegal(op.getFunctionType());
     });
 
     typeConverter.addConversion([](Type type) { return type; });
     typeConverter.addConversion(
         [](TupleType tupleType, SmallVectorImpl<Type> &types) {
           tupleType.getFlattenedTypes(types);
           return success();
         });
+    typeConverter.addArgumentMaterialization(buildMakeTupleOp);
 
-    decomposer.addDecomposeValueConversion([](OpBuilder &builder, Location loc,
-                                              TupleType resultType, Value value,
-                                              SmallVectorImpl<Value> &values) {
-      for (unsigned i = 0, e = resultType.size(); i < e; ++i) {
-        Value res = builder.create<test::GetTupleElementOp>(
-            loc, resultType.getType(i), value, builder.getI32IntegerAttr(i));
-        values.push_back(res);
-      }
-      return success();
-    });
-
-    typeConverter.addArgumentMaterialization(
-        [](OpBuilder &builder, TupleType resultType, ValueRange inputs,
-           Location loc) -> std::optional<Value> {
-          if (inputs.size() == 1)
-            return std::nullopt;
-          TupleType tuple = builder.getTupleType(inputs.getTypes());
-          Value value = builder.create<test::MakeTupleOp>(loc, tuple, inputs);
-          return value;
-        });
+    ValueDecomposer decomposer;
+    decomposer.addDecomposeValueConversion(buildDecomposeTuple);
 
     populateDecomposeCallGraphTypesPatterns(context, typeConverter, decomposer,
                                             patterns);
 
     if (failed(applyPartialConversion(module, target, std::move(patterns))))
       return signalPassFailure();
   }
 };