// 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();
// 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;
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++]);
}
}
-> VariableDecorator {
return VariableDecorator(cast<llvm::DefInit>(init)->getDef());
}
+
+auto tblgen::Operator::getArgToOperandOrAttribute(int index) const
+ -> OperandOrAttribute {
+ return attrOrOperandMapping[index];
+}
// 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";
// 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();
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);
}