[MLIR][TableGen] Fix ambiguous build methods when inferring result types.

[jurahul]

Summary

• Fix ODS framework to suppress build methods that infer result types and are ambiguous with collective variants. This applies to operations with a single variadic inputs whose result types can be inferred.
• Extended OpBuildGenTest to test these kinds of ops.

[stephenneuendorffer]

Seems like this will prevent the collision in the short term, but this solution doesn't seem like it will scale as builders get more complex. I'm not opposed to a special case here, but I'd feel more comfortable if there was a more general structure in place to prevent this from scaling to lots of special cases.

[jurahul]

I agree. A more general solution might be to generate all "collective" variants first, and then have a way for the build methods that will be ambiguous with any of them automatically filtered out, so that we don't have to guard their generation, just their emission. That might need the OpMethodSignature to have more structure than the current string members. I will file a bug for this assuming folks agree.

[stephenneuendorffer]

I'm happy with filing an issue about it.

[jurahul]

Thanks. @jpienaar, waiting to give you an opportunity to look as well.

[jpienaar]

In D85060#2191402, @jurahul wrote:

I agree. A more general solution might be to generate all "collective" variants first, and then have a way for the build methods that will be ambiguous with any of them automatically filtered out, so that we don't have to guard their generation, just their emission. That might need the OpMethodSignature to have more structure than the current string members. I will file a bug for this assuming folks agree.

Why not have opClass.newMethod flag this? If adding a new method to the class with the same signature, flag it (return nullptr or some such, depending on where the collision occurs it is either fine or an error). Start with the build methods users requested first and then add autogenerated ones as long as they don't conflict. Then we don't have to do gymnastics here. I'd say even the non-collective ones are the more user friendly ones. Although in cases of conflict it provides the same signature so perhaps it doesn't matter which one gets generated.

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

235

s/of/if/

[jurahul]

In D85060#2203065, @jpienaar wrote:

In D85060#2191402, @jurahul wrote:

I agree. A more general solution might be to generate all "collective" variants first, and then have a way for the build methods that will be ambiguous with any of them automatically filtered out, so that we don't have to guard their generation, just their emission. That might need the OpMethodSignature to have more structure than the current string members. I will file a bug for this assuming folks agree.

Why not have opClass.newMethod flag this? If adding a new method to the class with the same signature, flag it (return nullptr or some such, depending on where the collision occurs it is either fine or an error). Start with the build methods users requested first and then add autogenerated ones as long as they don't conflict. Then we don't have to do gymnastics here. I'd say even the non-collective ones are the more user friendly ones. Although in cases of conflict it provides the same signature so perhaps it doesn't matter which one gets generated.

That scheme would work as well. Although note ambiguity can come in 2 forms: (a) 2 different methods with exactly same signature (in which case both should be equivalent and generating any one of them is ok) and (b) 2 methods, one having a prefix of arguments as the other and the rest of the arguments in the other all have default values. In such a case, we want to preserve the second one (with more arguments). So ordering the generation of these methods will be important as well if we go with eager pruning. Otherwise, we just record all of them, and do pruning of the ambiguous ones at emission time.

[jurahul]

Fix comment

[jpienaar]

all have default values.

Oh yeah, that gets nasty in checking too. Almost makes me think of fluent style builders ;-)

LG to avoid current issues and we can refactor to avoid it in a more structured way post.

[jurahul]

Thanks. Filed https://bugs.llvm.org/show_bug.cgi?id=47095 to tract the more general mechanism for doing this.

mlir/include/mlir/TableGen/Operator.h

public:
   }
 
   // Returns the number of variadic operands in this operation.
   unsigned getNumVariableLengthOperands() const;
 
   // Returns the total number of arguments.
   int getNumArgs() const { return arguments.size(); }
 
+  // Returns true of the operation has a single variadic arg.
+  bool hasSingleVariadicArg() const;
+
+  // Returns true if the operation has a single variadic result.
+  bool hasSingleVariadicResult() const {
+    return getNumResults() == 1 && getResult(0).isVariadic();
+  }
+
+  // Returns true of the operation has no variadic regions.
+  bool hasNoVariadicRegions() const { return getNumVariadicRegions() == 0; }
+
   using arg_iterator = const Argument *;
   using arg_range = llvm::iterator_range<arg_iterator>;
 
   // Op argument (attribute or operand) iterators.
   arg_iterator arg_begin() const;
   arg_iterator arg_end() const;
   arg_range getArgs() const;
 

mlir/lib/TableGen/Operator.cpp

 }
 
 unsigned tblgen::Operator::getNumVariableLengthOperands() const {
   return llvm::count_if(operands, [](const NamedTypeConstraint &c) {
     return c.constraint.isVariableLength();
   });
 }
 
+bool tblgen::Operator::hasSingleVariadicArg() const {
+  return getNumArgs() == 1 && getArg(0).is<tblgen::NamedTypeConstraint *>() &&
+         getOperand(0).isVariadic();
+}
+
 tblgen::Operator::arg_iterator tblgen::Operator::arg_begin() const {
   return arguments.begin();
 }
 
 tblgen::Operator::arg_iterator tblgen::Operator::arg_end() const {
   return arguments.end();
 }
 

mlir/test/lib/Dialect/Test/TestOps.td

 }
 
 // Single variadic arg and multiple variadic results.
 def TableGenBuildOp3 : TEST_Op<"tblgen_build_3", [SameVariadicResultSize]> {
   let arguments = (ins Variadic<AnyType>:$inputs);
   let results = (outs Variadic<AnyType>:$resultA, Variadic<AnyType>:$resultB);
 }
 
+// Single variadic arg, non variadic results, with SameOperandsAndResultType.
+// Tests suppression of ambiguious build methods for operations with
+// SameOperandsAndResultType trait.
+def TableGenBuildOp4 : TEST_Op<"tblgen_build_4", [SameOperandsAndResultType]> {
+  let arguments = (ins Variadic<AnyType>:$inputs);
+  let results = (outs AnyType:$result);
+}
+
+// Single variadic arg with SameOperandsAndResultType and InferTypeOpInterface.
+// Tests suppression of ambiguious build methods for operations with
+// SameOperandsAndResultType and InferTypeOpInterface.
+def TableGenBuildOp5 : TEST_Op<"tblgen_build_5",
+      [SameOperandsAndResultType, InferTypeOpInterface]> {
+  let arguments = (ins Variadic<AnyType>:$inputs);
+  let results = (outs AnyType:$result);
+
+  let extraClassDeclaration = [{
+    static LogicalResult inferReturnTypes(MLIRContext *, 
+          Optional<Location> location, ValueRange operands,
+          DictionaryAttr attributes, RegionRange regions,
+          SmallVectorImpl<Type> &inferredReturnTypes) {
+      inferredReturnTypes.assign({operands[0].getType()});
+      return success();
+    }
+   }];
+}
+
 #endif // TEST_OPS

mlir/test/mlir-tblgen/op-result.td

 // Test that if the only operand is variadic, we access the first value in the
 // pack to set result type
 // ---
 def OpK : NS_Op<"only_input_is_variadic_with_same_value_type_op", [SameOperandsAndResultType]> {
   let arguments = (ins Variadic<AnyTensor>:$input);
   let results = (outs AnyTensor:$result);
 }
 
-// CHECK-LABEL: OpK::build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::ValueRange input)
-// CHECK: odsState.addTypes({input.front().getType()});
+// CHECK-LABEL: OpK::build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes )
+// CHECK: odsState.addTypes({operands[0].getType()});
 
 // Test with inferred shapes and interleaved with operands/attributes.
 //
 def OpL : NS_Op<"op_with_all_types_constraint",
     [AllTypesMatch<["a", "b"]>]> {
   let arguments = (ins I32Attr:$attr1, AnyType:$a);
   let results = (outs Res<AnyType, "output b", []>:$b);
 }
 
 // CHECK-LABEL: LogicalResult OpL::inferReturnTypes
 // CHECK-NOT: }
 // CHECK: inferredReturnTypes[0] = operands[0].getType();

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

private:
   // type as all results' types.
   void genUseOperandAsResultTypeSeparateParamBuilder();
 
   // Generates the build() method that takes all operands/attributes
   // collectively as one parameter. The generated build() method uses first
   // operand's type as all results' types.
   void genUseOperandAsResultTypeCollectiveParamBuilder();
 
+  // Returns true if the inferred collective param build method should be
+  // generated.
+  bool shouldGenerateInferredTypeCollectiveParamBuilder();
+
   // Generates the build() method that takes aggregate operands/attributes
   // parameters. This build() method uses inferred types as result types.
   // Requires: The type needs to be inferable via InferTypeOpInterface.
   void genInferredTypeCollectiveParamBuilder();
 
   // Generates the build() method that takes each operand/attribute as a
   // stand-alone parameter. The generated build() method uses first attribute's
   // type as all result's types.
void OpEmitter::genSeparateArgParamBuilder() {
   //    in the build methods generated here) and the single arg must be a
   //    non-attribute variadic argument.
   // 3. [single `ArrayReg<Type>` parameter for result types]:
   //      3a. paramKind should be Collective, or
   //      3b. paramKind should be Separate and there should be a single variadic
   //          result
   //
   // In that case, skip generating such ambiguous build methods here.
-  bool hasSingleVariadicResult =
-      op.getNumResults() == 1 && op.getResult(0).isVariadic();
-
-  bool hasSingleVariadicArg =
-      op.getNumArgs() == 1 &&
-      op.getArg(0).is<tblgen::NamedTypeConstraint *>() &&
-      op.getOperand(0).isVariadic();
-  bool hasNoVariadicRegions = op.getNumVariadicRegions() == 0;
-
   for (auto attrType : attrBuilderType) {
     // Case 3b above.
-    if (!(hasNoVariadicRegions && hasSingleVariadicArg &&
-          hasSingleVariadicResult))
+    if (!(op.hasNoVariadicRegions() && op.hasSingleVariadicArg() &&
+          op.hasSingleVariadicResult()))
       emit(attrType, TypeParamKind::Separate, /*inferType=*/false);
-    if (canInferType(op))
+    if (canInferType(op)) {
+      // When inferType = true, the generated build method does not have
+      // result types. If the op has a single variadic arg, then this build
+      // method will be ambiguious with the collective inferred build method
+      // generated in `genInferredTypeCollectiveParamBuilder`. If we are going
+      // to generate that collective inferred method, suppress generating the
+      // ambiguious build method here.
+      bool buildMethodAmbiguious =
+          op.hasSingleVariadicArg() &&
+          shouldGenerateInferredTypeCollectiveParamBuilder();
+      if (!buildMethodAmbiguious)
         emit(attrType, TypeParamKind::None, /*inferType=*/true);
+    }
     // The separate arg + collective param kind method will be:
     // (a) Same as the separate arg + separate param kind method if there is
     //     only one variadic result.
     // (b) Ambiguous with the collective params method under conditions in (3a)
     //     above.
     // In either case, skip generating such build method.
-    if (!hasSingleVariadicResult &&
-        !(hasNoVariadicRegions && hasSingleVariadicArg))
+    if (!op.hasSingleVariadicResult() &&
+        !(op.hasNoVariadicRegions() && op.hasSingleVariadicArg()))
       emit(attrType, TypeParamKind::Collective, /*inferType=*/false);
   }
 }
 
 void OpEmitter::genUseOperandAsResultTypeCollectiveParamBuilder() {
-  // If this op has a variadic result, we cannot generate this builder because
-  // we don't know how many results to create.
-  if (op.getNumVariableLengthResults() != 0)
-    return;
-
   int numResults = op.getNumResults();
 
   // Signature
   std::string params =
       std::string("::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &") +
       builderOpState +
       ", ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> "
       "attributes";
void OpEmitter::genUseOperandAsResultTypeCollectiveParamBuilder() {
   }
 
   // Result types
   SmallVector<std::string, 2> resultTypes(numResults, "operands[0].getType()");
   body << "  " << builderOpState << ".addTypes({"
        << llvm::join(resultTypes, ", ") << "});\n\n";
 }
 
+bool OpEmitter::shouldGenerateInferredTypeCollectiveParamBuilder() {
+  return canInferType(op) && op.getNumSuccessors() == 0;
+}
+
 void OpEmitter::genInferredTypeCollectiveParamBuilder() {
   // TODO: Expand to support regions.
   std::string params =
       std::string("::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &") +
       builderOpState +
       ", ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> "
       "attributes = {}";
   auto &m = opClass.newMethod("void", "build", params, OpMethod::MP_Static);
void OpEmitter::genBuilder() {
   // 2. one having an aggregated parameter for all result types / operands /
   //    attributes, and
   genCollectiveParamBuilder();
   // 3. one having a stand-alone parameter for each operand and attribute,
   //    use the first operand or attribute's type as all result types
   //    to facilitate different call patterns.
   if (op.getNumVariableLengthResults() == 0) {
     if (op.getTrait("OpTrait::SameOperandsAndResultType")) {
+      // If the operation has a single variadic input, then the build method
+      // generated by `genUseOperandAsResultTypeSeparateParamBuilder` will be
+      // ambiguious with the one generated by
+      // `genUseOperandAsResultTypeCollectiveParamBuilder` (they both will have
+      // a single `ValueRange` argument for operands, and the collective one
+      // will have a `ArrayRef<NamedAttribute>` argument initalized to empty).
+      // Suppress such ambiguious build method.
+      if (!op.hasSingleVariadicArg())
         genUseOperandAsResultTypeSeparateParamBuilder();
+
+      // The build method generated by the inferred type collective param
+      // builder and one generated here have the same arguments and hence
+      // generating both will be ambiguious. Enable just one of them.
+      if (!shouldGenerateInferredTypeCollectiveParamBuilder())
         genUseOperandAsResultTypeCollectiveParamBuilder();
     }
     if (op.getTrait("OpTrait::FirstAttrDerivedResultType"))
       genUseAttrAsResultTypeBuilder();
   }
 }
 
 void OpEmitter::genCollectiveParamBuilder() {
   int numResults = op.getNumResults();
void OpEmitter::genCollectiveParamBuilder() {
   if (numVariadicResults == 0 || numNonVariadicResults != 0)
     body << "  assert(resultTypes.size()"
          << (numVariadicResults != 0 ? " >= " : " == ") << numNonVariadicResults
          << "u && \"mismatched number of return types\");\n";
   body << "  " << builderOpState << ".addTypes(resultTypes);\n";
 
   // Generate builder that infers type too.
   // TODO: Expand to handle regions and successors.
-  if (canInferType(op) && op.getNumSuccessors() == 0)
+  if (shouldGenerateInferredTypeCollectiveParamBuilder())
     genInferredTypeCollectiveParamBuilder();
 }
 
 void OpEmitter::buildParamList(std::string &paramList,
                                SmallVectorImpl<std::string> &resultTypeNames,
                                TypeParamKind typeParamKind,
                                AttrParamKind attrParamKind) {
   resultTypeNames.clear();

mlir/unittests/TableGen/OpBuildGen.cpp

void verifyOp(OpTy &&concreteOp, std::vector<Type> resultTypes,
 
     EXPECT_EQ(op->getAttrs().size(), attrs.size());
     for (unsigned idx : llvm::seq<unsigned>(0U, attrs.size()))
       EXPECT_EQ(op->getAttr(attrs[idx].first.strref()), attrs[idx].second);
 
     concreteOp.erase();
   }
 
+  // Helper method to test ops with inferred result types and single variadic
+  // input.
+  template <typename OpTy>
+  void testSingleVariadicInputInferredType() {
+    // Test separate arg, separate param build method.
+    auto op = builder.create<OpTy>(loc, i32Ty, ArrayRef<Value>{cstI32, cstI32});
+    verifyOp(std::move(op), {i32Ty}, {cstI32, cstI32}, noAttrs);
+
+    // Test collective params build method.
+    op = builder.create<OpTy>(loc, ArrayRef<Type>{i32Ty},
+                              ArrayRef<Value>{cstI32, cstI32});
+    verifyOp(std::move(op), {i32Ty}, {cstI32, cstI32}, noAttrs);
+
+    // Test build method with no result types, default value of attributes.
+    op = builder.create<OpTy>(loc, ArrayRef<Value>{cstI32, cstI32});
+    verifyOp(std::move(op), {i32Ty}, {cstI32, cstI32}, noAttrs);
+
+    // Test build method with no result types and supplied attributes.
+    op = builder.create<OpTy>(loc, ArrayRef<Value>{cstI32, cstI32}, attrs);
+    verifyOp(std::move(op), {i32Ty}, {cstI32, cstI32}, attrs);
+  }
+
 protected:
   MLIRContext ctx;
   OpBuilder builder;
   Location loc;
   Type i32Ty;
   Type f32Ty;
   TableGenConstant cstI32;
   TableGenConstant cstF32;
TEST_F(OpBuildGenTest,
   verifyOp(std::move(op), {i32Ty, f32Ty}, {cstI32}, noAttrs);
 
   // Test collective params build method, with attributes.
   op = builder.create<TableGenBuildOp3>(loc, ArrayRef<Type>{i32Ty, f32Ty},
                                         ArrayRef<Value>{cstI32}, attrs);
   verifyOp(std::move(op), {i32Ty, f32Ty}, {cstI32}, attrs);
 }
 
+// The next 2 tests test supression of ambiguious build methods for ops that
+// have a single variadic input, and single non-variadic result, and which
+// support the SameOperandsAndResultType trait and and optionally the
+// InferOpTypeInterface interface. For such ops, the ODS framework generates
+// build methods with no result types as they are inferred from the input types.
+TEST_F(OpBuildGenTest, BuildMethodsSameOperandsAndResultTypeSuppression) {
+  testSingleVariadicInputInferredType<TableGenBuildOp4>();
+}
+
+TEST_F(
+    OpBuildGenTest,
+    BuildMethodsSameOperandsAndResultTypeAndInferOpTypeInterfaceSuppression) {
+  testSingleVariadicInputInferredType<TableGenBuildOp5>();
+}
+
 } // namespace mlir