xref: /llvm-project/mlir/lib/Dialect/IRDL/IR/IRDLOps.cpp (revision b86a9c5bf2fab0408a3d549995d6e2449f71a16d)
1 //===- IRDLOps.cpp - IRDL dialect -------------------------------*- C++ -*-===//
2 //
3 // This file is licensed under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/Dialect/IRDL/IR/IRDL.h"
10 #include "mlir/Dialect/IRDL/IRDLSymbols.h"
11 #include "mlir/IR/ValueRange.h"
12 #include <optional>
13 
14 using namespace mlir;
15 using namespace mlir::irdl;
16 
17 /// Maps given `args` to the index in the `valueToConstr`
18 static SmallVector<unsigned>
getConstraintIndicesForArgs(mlir::OperandRange args,ArrayRef<Value> valueToConstr)19 getConstraintIndicesForArgs(mlir::OperandRange args,
20                             ArrayRef<Value> valueToConstr) {
21   SmallVector<unsigned> constraints;
22   for (Value arg : args) {
23     for (auto [i, value] : enumerate(valueToConstr)) {
24       if (value == arg) {
25         constraints.push_back(i);
26         break;
27       }
28     }
29   }
30   return constraints;
31 }
32 
getVerifier(ArrayRef<Value> valueToConstr,DenseMap<TypeOp,std::unique_ptr<DynamicTypeDefinition>> const & types,DenseMap<AttributeOp,std::unique_ptr<DynamicAttrDefinition>> const & attrs)33 std::unique_ptr<Constraint> IsOp::getVerifier(
34     ArrayRef<Value> valueToConstr,
35     DenseMap<TypeOp, std::unique_ptr<DynamicTypeDefinition>> const &types,
36     DenseMap<AttributeOp, std::unique_ptr<DynamicAttrDefinition>> const
37         &attrs) {
38   return std::make_unique<IsConstraint>(getExpectedAttr());
39 }
40 
getVerifier(ArrayRef<Value> valueToConstr,DenseMap<TypeOp,std::unique_ptr<DynamicTypeDefinition>> const & types,DenseMap<AttributeOp,std::unique_ptr<DynamicAttrDefinition>> const & attrs)41 std::unique_ptr<Constraint> BaseOp::getVerifier(
42     ArrayRef<Value> valueToConstr,
43     DenseMap<TypeOp, std::unique_ptr<DynamicTypeDefinition>> const &types,
44     DenseMap<AttributeOp, std::unique_ptr<DynamicAttrDefinition>> const
45         &attrs) {
46   MLIRContext *ctx = getContext();
47 
48   // Case where the input is a symbol reference.
49   // This corresponds to the case where the base is an IRDL type or attribute.
50   if (auto baseRef = getBaseRef()) {
51     // The verifier for BaseOp guarantees it is within a dialect.
52     Operation *defOp =
53         irdl::lookupSymbolNearDialect(getOperation(), baseRef.value());
54 
55     // Type case.
56     if (auto typeOp = dyn_cast<TypeOp>(defOp)) {
57       DynamicTypeDefinition *typeDef = types.at(typeOp).get();
58       auto name = StringAttr::get(ctx, typeDef->getDialect()->getNamespace() +
59                                            "." + typeDef->getName().str());
60       return std::make_unique<BaseTypeConstraint>(typeDef->getTypeID(), name);
61     }
62 
63     // Attribute case.
64     auto attrOp = cast<AttributeOp>(defOp);
65     DynamicAttrDefinition *attrDef = attrs.at(attrOp).get();
66     auto name = StringAttr::get(ctx, attrDef->getDialect()->getNamespace() +
67                                          "." + attrDef->getName().str());
68     return std::make_unique<BaseAttrConstraint>(attrDef->getTypeID(), name);
69   }
70 
71   // Case where the input is string literal.
72   // This corresponds to the case where the base is a registered type or
73   // attribute.
74   StringRef baseName = getBaseName().value();
75 
76   // Type case.
77   if (baseName[0] == '!') {
78     auto abstractType = AbstractType::lookup(baseName.drop_front(1), ctx);
79     if (!abstractType) {
80       emitError() << "no registered type with name " << baseName;
81       return nullptr;
82     }
83     return std::make_unique<BaseTypeConstraint>(abstractType->get().getTypeID(),
84                                                 abstractType->get().getName());
85   }
86 
87   auto abstractAttr = AbstractAttribute::lookup(baseName.drop_front(1), ctx);
88   if (!abstractAttr) {
89     emitError() << "no registered attribute with name " << baseName;
90     return nullptr;
91   }
92   return std::make_unique<BaseAttrConstraint>(abstractAttr->get().getTypeID(),
93                                               abstractAttr->get().getName());
94 }
95 
getVerifier(ArrayRef<Value> valueToConstr,DenseMap<TypeOp,std::unique_ptr<DynamicTypeDefinition>> const & types,DenseMap<AttributeOp,std::unique_ptr<DynamicAttrDefinition>> const & attrs)96 std::unique_ptr<Constraint> ParametricOp::getVerifier(
97     ArrayRef<Value> valueToConstr,
98     DenseMap<TypeOp, std::unique_ptr<DynamicTypeDefinition>> const &types,
99     DenseMap<AttributeOp, std::unique_ptr<DynamicAttrDefinition>> const
100         &attrs) {
101   SmallVector<unsigned> constraints =
102       getConstraintIndicesForArgs(getArgs(), valueToConstr);
103 
104   // Symbol reference case for the base.
105   // The verifier for ParametricOp guarantees it is within a dialect.
106   SymbolRefAttr symRef = getBaseType();
107   Operation *defOp = irdl::lookupSymbolNearDialect(getOperation(), symRef);
108   if (!defOp) {
109     emitError() << symRef << " does not refer to any existing symbol";
110     return nullptr;
111   }
112 
113   if (auto typeOp = dyn_cast<TypeOp>(defOp))
114     return std::make_unique<DynParametricTypeConstraint>(types.at(typeOp).get(),
115                                                          constraints);
116 
117   if (auto attrOp = dyn_cast<AttributeOp>(defOp))
118     return std::make_unique<DynParametricAttrConstraint>(attrs.at(attrOp).get(),
119                                                          constraints);
120 
121   llvm_unreachable("verifier should ensure that the referenced operation is "
122                    "either a type or an attribute definition");
123 }
124 
getVerifier(ArrayRef<Value> valueToConstr,DenseMap<TypeOp,std::unique_ptr<DynamicTypeDefinition>> const & types,DenseMap<AttributeOp,std::unique_ptr<DynamicAttrDefinition>> const & attrs)125 std::unique_ptr<Constraint> AnyOfOp::getVerifier(
126     ArrayRef<Value> valueToConstr,
127     DenseMap<TypeOp, std::unique_ptr<DynamicTypeDefinition>> const &types,
128     DenseMap<AttributeOp, std::unique_ptr<DynamicAttrDefinition>> const
129         &attrs) {
130   return std::make_unique<AnyOfConstraint>(
131       getConstraintIndicesForArgs(getArgs(), valueToConstr));
132 }
133 
getVerifier(ArrayRef<Value> valueToConstr,DenseMap<TypeOp,std::unique_ptr<DynamicTypeDefinition>> const & types,DenseMap<AttributeOp,std::unique_ptr<DynamicAttrDefinition>> const & attrs)134 std::unique_ptr<Constraint> AllOfOp::getVerifier(
135     ArrayRef<Value> valueToConstr,
136     DenseMap<TypeOp, std::unique_ptr<DynamicTypeDefinition>> const &types,
137     DenseMap<AttributeOp, std::unique_ptr<DynamicAttrDefinition>> const
138         &attrs) {
139   return std::make_unique<AllOfConstraint>(
140       getConstraintIndicesForArgs(getArgs(), valueToConstr));
141 }
142 
getVerifier(ArrayRef<Value> valueToConstr,DenseMap<TypeOp,std::unique_ptr<DynamicTypeDefinition>> const & types,DenseMap<AttributeOp,std::unique_ptr<DynamicAttrDefinition>> const & attrs)143 std::unique_ptr<Constraint> AnyOp::getVerifier(
144     ArrayRef<Value> valueToConstr,
145     DenseMap<TypeOp, std::unique_ptr<DynamicTypeDefinition>> const &types,
146     DenseMap<AttributeOp, std::unique_ptr<DynamicAttrDefinition>> const
147         &attrs) {
148   return std::make_unique<AnyAttributeConstraint>();
149 }
150 
getVerifier(ArrayRef<Value> valueToConstr,DenseMap<TypeOp,std::unique_ptr<DynamicTypeDefinition>> const & types,DenseMap<AttributeOp,std::unique_ptr<DynamicAttrDefinition>> const & attrs)151 std::unique_ptr<RegionConstraint> RegionOp::getVerifier(
152     ArrayRef<Value> valueToConstr,
153     DenseMap<TypeOp, std::unique_ptr<DynamicTypeDefinition>> const &types,
154     DenseMap<AttributeOp, std::unique_ptr<DynamicAttrDefinition>> const
155         &attrs) {
156   return std::make_unique<RegionConstraint>(
157       getConstrainedArguments() ? std::optional{getConstraintIndicesForArgs(
158                                       getEntryBlockArgs(), valueToConstr)}
159                                 : std::nullopt,
160       getNumberOfBlocks());
161 }
162