[ods] Update Operator to record Arg->[Attr|Operand]Index mapping

[jpienaar]

Summary

Also fixed bug in type inferface generator to address bug where operands and
attributes are interleaved.

[rriddle]

mlir/include/mlir/TableGen/Operator.h

266

nit: Can you bitpack these together?

270

nit: argument at index?

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1605

nit: Drop the mir::.

[jpienaar]

Packed to one int with member accessors.

mlir/include/mlir/TableGen/Operator.h

private:
     Optional<int> index;
     Optional<TypeConstraint> constraint;
   };
 
   // Return all arguments or type constraints with same type as result[index].
   // Requires: all result types are known.
   ArrayRef<ArgOrType> getSameTypeAsResult(int index) const;
 
+  // Pair consisting kind of argument and index into operands or attributes.
+  struct OperandOrAttribute {
+    enum class Kind { Operand, Attribute };
+    OperandOrAttribute(Kind kind, int index) {
+      packed = (index << 1) & (kind == Kind::Attribute);
+    }
+    int operandOrAttributeIndex() const { return (packed >> 1); }
+    Kind kind() { return (packed & 0x1) ? Kind::Attribute : Kind::Operand; }
+
+  private:
+    int packed;
+  };
+
+  // Returns the OperandOrAttribute corresponding to the index.
+  OperandOrAttribute getArgToOperandOrAttribute(int index) const;
+
 private:
   // Populates the vectors containing operands, attributes, results and traits.
   void populateOpStructure();
 
   // Populates type inference info (mostly equality) with input a mapping from
   // names to indices for arguments and results.
   void populateTypeInferenceInfo(
       const llvm::StringMap<int> &argumentsAndResultsIndex);
private:
   SmallVector<OpTrait, 4> traits;
 
   // The regions of this op.
   SmallVector<NamedRegion, 1> regions;
 
   // The argument with the same type as the result.
   SmallVector<SmallVector<ArgOrType, 2>, 4> resultTypeMapping;
 
+  // Map from argument to attribute or operand number.
+  SmallVector<OperandOrAttribute, 4> attrOrOperandMapping;
+
   // The number of native attributes stored in the leading positions of
   // `attributes`.
   int numNativeAttributes;
 
   // The TableGen definition of this op.
   const llvm::Record &def;
 
   // Whether the type of all results are known.
   bool allResultsHaveKnownTypes;
 };
 
 } // end namespace tblgen
 } // end namespace mlir
 
 #endif // MLIR_TABLEGEN_OPERATOR_H_

mlir/lib/TableGen/Operator.cpp

void tblgen::Operator::populateOpStructure() {
   // elements.
   int operandIndex = 0, attrIndex = 0;
   for (unsigned i = 0; i != numArgs; ++i) {
     Record *argDef = dyn_cast<DefInit>(argumentValues->getArg(i))->getDef();
     if (argDef->isSubClassOf(opVarClass))
       argDef = argDef->getValueAsDef("constraint");
 
     if (argDef->isSubClassOf(typeConstraintClass)) {
+      attrOrOperandMapping.push_back(
+          {OperandOrAttribute::Kind::Operand, operandIndex});
       arguments.emplace_back(&operands[operandIndex++]);
     } else {
       assert(argDef->isSubClassOf(attrClass));
+      attrOrOperandMapping.push_back(
+          {OperandOrAttribute::Kind::Attribute, attrIndex});
       arguments.emplace_back(&attributes[attrIndex++]);
     }
   }
 
   auto *resultsDag = def.getValueAsDag("results");
   auto *outsOp = dyn_cast<DefInit>(resultsDag->getOperator());
   if (!outsOp || outsOp->getDef()->getName() != "outs") {
     PrintFatalError(def.getLoc(), "'results' must have 'outs' directive");
else
       os << "[operand] " << arg.get<NamedTypeConstraint *>()->name << '\n';
   }
 }
 
 auto tblgen::Operator::VariableDecoratorIterator::unwrap(llvm::Init *init)
     -> VariableDecorator {
   return VariableDecorator(cast<llvm::DefInit>(init)->getDef());
 }
+
+auto tblgen::Operator::getArgToOperandOrAttribute(int index) const
+    -> OperandOrAttribute {
+  return attrOrOperandMapping[index];
+}

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

 // RUN: mlir-tblgen -gen-op-defs -I %S/../../include %s | FileCheck %s
 
 include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 
 def Test_Dialect : Dialect {
   let name = "test";
 }
 class NS_Op<string mnemonic, list<OpTrait> traits> :
     Op<Test_Dialect, mnemonic, traits>;
 
 def OpA : NS_Op<"one_normal_result_op", []> {

 // ---
 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()});
+
+// 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

void OpEmitter::genTypeInterfaceMethods() {
   fctx.withBuilder("odsBuilder");
   os << "  ::mlir::Builder odsBuilder(context);\n";
 
   auto emitType =
       [&](const tblgen::Operator::ArgOrType &type) -> OpMethodBody & {
     if (type.isArg()) {
       auto argIndex = type.getArg();
       assert(!op.getArg(argIndex).is<NamedAttribute *>());
-      return os << "operands[" << argIndex << "].getType()";
+      auto arg = op.getArgToOperandOrAttribute(argIndex);
+      if (arg.kind() == Operator::OperandOrAttribute::Kind::Operand)
+        return os << "operands[" << arg.operandOrAttributeIndex()
+                  << "].getType()";
+      return os << "attributes[" << arg.operandOrAttributeIndex()
+                << "].getType()";
     } else {
       return os << tgfmt(*type.getType().getBuilderCall(), &fctx);
     }
   };
 
   for (int i = 0, e = op.getNumResults(); i != e; ++i) {
     os << "  inferredReturnTypes[" << i << "] = ";
     auto types = op.getSameTypeAsResult(i);