[mlir][linalg] Verify indexing map required attributes

[antiagainst]

Summary

Indexing maps for named ops can reference attributes so that
we can synthesize the indexing map dynamically. This supports
cases like strides for convolution ops. However, it does cause
an issue: now the indexing_maps() function call is dependent
on those attributes.

Linalg ops inherit LinalgOpInterfaceTraits, which calls
verifyStructuredOpInterface() to verify the interface.
verifyStructuredOpInterface() further calls indexing_maps().
Note that trait verification is done before the op itself,
where ODS generates the verification for those attributes.
So we can have indexing_maps() referencing non-existing or
invalid attribute, before the ODS-generated verification
kick in.

There isn't a dependency handling mechansim for traits.
This commit adds new interface methods to query whether an
op hasDynamicIndexingMaps() and then perform
verifyIndexingMapRequiredAttributes() in
verifyStructuredOpInterface() to handle the dependency issue.

[mehdi_amini]

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

307

Do we need this check? Shouldn't we be able to call the verify unconditionally?

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

1849

I think it is a good practice for format string with more than one parameter to document their list of parameters, this one has 7 now :)

2067

I think what's happening in this method could benefit from more documentation

[nicolasvasilache]

Approving conditioned on the extra doc @mehdi_amini asked for. Thanks!

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

307

Not all ops have dynamic indexing maps (i.e. ones that are dependent on the op instance's attributes).
A particular op knows whether its definition requires such attributes to define the indexing maps.
So this all makes sense to me and we are calling the verifier unconditionally.

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

1849

+1, some /*stringName=*/ would help read this better.

[antiagainst]

• Address comments

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

307

Yeah. I feel this is more clear logically to show that not every op need this kind of check given that not all of them has dynamic indexing maps. :)

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

1849

Added explanation ahead .

[mehdi_amini]

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

307

I mean that calling verifyIndexingMapRequiredAttributes() when hasDynamicIndexingMaps() return false will succeed.
So you should either make it fails when hasDynamicIndexingMaps() returns false, or not guard it here.

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

InterfaceMethod<
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.iterator_types();
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Return true if the indexing map is depending on the current op instance.
+        This means that the indexing map is dynamically synthesized by using the
+        op instance's concrete attributes, instead of being static for all
+        instances of the same op kind.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"hasDynamicIndexingMaps",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{ return false; }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Verify all attributes used by indexing maps are valid.
+      }],
+      /*retTy=*/"LogicalResult",
+      /*methodName=*/"verifyIndexingMapRequiredAttributes",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{ return success(); }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the indexing maps attribute within the current operation.
       }],
       /*retTy=*/"ArrayAttr",
       /*methodName=*/"indexing_maps"
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the indexing maps within the current operation.

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

if (nShapedOperands == 0)
     return linalgOp.emitOpError("expected at least 1 Shaped operand");
   if (failed(OpTrait::impl::verifyAtLeastNOperands(op, nShapedOperands)))
     return failure();
   // Should have at least one output tensor per result tensor.
   // Can also have outbut buffers that do not correspond to results.
   if (op->getNumResults() > linalgOp.getNumOutputTensors())
     return op->emitError("unexpected #results > #outputs");
 
+  // Before checking indexing maps, we need to make sure the attributes
+  // referenced by it are valid.
+  if (linalgOp.hasDynamicIndexingMaps())
+    if (failed(linalgOp.verifyIndexingMapRequiredAttributes()))
+      return failure();
+
   // All shaped operands must be indexed.
   if (linalgOp.indexing_maps().size() != linalgOp.getNumShapedOperands())
     return linalgOp.emitOpError("expected the number of indexing_map (")
            << linalgOp.indexing_maps().size()
            << ") to be equal to the number of shaped operands ("
            << linalgOp.getNumShapedOperands() << ")";
 
   SmallVector<AffineMap, 4> indexingMaps;

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc

 // ODS: F32ArrayAttr:$array_attr,
 // ODS: F32:$f32_attr,
 // ODS: RankedF32ElementsAttr<[4]>:$fvec_attr,
 // ODS: I32:$i32_attr,
 // ODS: I64:$i64_attr,
 // ODS: RankedI32ElementsAttr<[5, 6]>:$ivec_attr,
 // ODS: OptionalAttr<F32>:$optional_attr
 //
+// ODS: bool hasDynamicIndexingMaps();
+// ODS: LogicalResult verifyIndexingMapRequiredAttributes();
+//
+// IMPL: bool Test4Op::hasDynamicIndexingMaps() { return true; }
+// IMPL: LogicalResult Test4Op::verifyIndexingMapRequiredAttributes()
+// IMPL:   op->getAttrOfType<ArrayAttr>("array_attr")
+// IMPL:   op->getAttr("f32_attr")
+// IMPL:   op->getAttrOfType<DenseElementsAttr>("fvec_attr")
+// IMPL:   op->getAttr("i32_attr")
+// IMPL:   op->getAttr("i64_attr")
+// IMPL:   op->getAttrOfType<DenseElementsAttr>("ivec_attr")
+//
 ods_def<Test4Op> :
 def test4(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N))
 attr(
   f32_attr: f32,
   i32_attr: i32,
   i64_attr: i64,
   fvec_attr: 4xf32,
   ivec_attr: 5x6xi32,

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

public:
   void printODS(llvm::raw_ostream &os, StringRef cppOpName,
                 StringRef linalgOpName, ArrayRef<StringRef> interfaces,
                 ComprehensionParsingState &state);
 
   /// Print the C++ StructuredOpsInterface impl of `iterator_types`.
   void printReferenceIterators(llvm::raw_ostream &os, StringRef cppOpName,
                                ComprehensionParsingState &state);
 
+  /// Print methods related to indexing map required attributes.
+  ///
+  /// Specifically, this prints the definitions for the following methods:
+  ///   bool hasDynamicIndexingMaps();
+  ///   LogicalResult verifyIndexingMapRequiredAttributes();
+  void printIndexingMapRequiredAttrMethods(llvm::raw_ostream &os,
+                                           StringRef cppOpName,
+                                           ComprehensionParsingState &state);
+
   /// Print the C++ StructuredOpsInterface impl of `indexing_maps`.
   void printReferenceIndexingMaps(llvm::raw_ostream &os, StringRef cppOpName,
                                   ComprehensionParsingState &state);
 
   /// Print the C++ StructuredOpsInterface impl of `regionBuilder`.
   void printRegionBuilder(llvm::raw_ostream &os, StringRef cppOpName,
                           ComprehensionParsingState &state);
 
if (genODSDecl) {
     printODS(os, cppOpName, tcName, interfaces, state);
     os << "\n";
   }
   if (genODSImpl) {
     auto &state = perComprehensionStates.back();
     std::string extraMethods;
     llvm::raw_string_ostream ss(extraMethods);
     printReferenceIterators(ss, cppOpName, state);
+    printIndexingMapRequiredAttrMethods(ss, cppOpName, state);
     printReferenceIndexingMaps(ss, cppOpName, state);
     printRegionBuilder(ss, cppOpName, state);
     printCanonicalizersAndFolders(ss, cppOpName);
     ss.flush();
     os << extraMethods << "\n";
   }
 
   return success();
for (const auto &attr : registeredAttrs) {
 
     attributes.push_back(llvm::formatv("{0}:${1}", odsType, name));
   }
 
   std::string attrList = llvm::join(attributes, ",\n");
   if (!attrList.empty())
     attrList = ",\n" + attrList;
 
+  // Template for Linalg named ops' ODS definitions. Parameters:
+  // {0}: ODS/C++ op name
+  // {1}: assembly op mnemonic
+  // {2}: op interface list
+  // {3}: documentation (summary + description)
+  // {4}: op attribute list
+  // {5}: the number of arguments for the op region
+  // {6}: builder methods taking standalone attribute parameters
+  // {7}: additional methods for attributes used by indexing maps
   const char *header = R"FMT(  def {0} : LinalgStructuredBase_Op<"{1}", [
     AttrSizedOperandSegments,
     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
     SingleBlockImplicitTerminator<"YieldOp">
     /*extraInterfaces=*/{2}]> {
       {3}
       let arguments = (ins
         Variadic<AnyShaped>:$inputs,
         Variadic<AnyShaped>:$outputs{4}
/*extraInterfaces=*/{2}]> {
           return regionBuilder;
         }
 
         // Generic methods.
         static unsigned getNumRegionArgs() {{ return {5}; }
         std::string getLibraryCallName() {{
           return generateLibraryCallName(getOperation());
         }
+
+        {7}
       }];
   })FMT";
 
   // Generate the list of extra implemented interfaces.
   std::string interfaceNameList;
   if (!interfaces.empty()) {
     llvm::raw_string_ostream ss(interfaceNameList);
     ss << ", "; // Leading comma to concat to existing list of interfaces.
const char *builderFmt = R"FMT(
           TypeRange(outputs));
         {2}
       }]>
     )FMT";
     attrBuilder =
         llvm::formatv(builderFmt, cppOpName, attrParamsList, attrStmtsList);
   }
 
+  std::string attrMethods;
+  if (!registeredAttrs.empty()) {
+    attrMethods = R"(
+      bool hasDynamicIndexingMaps();
+      LogicalResult verifyIndexingMapRequiredAttributes();
+    )";
+  }
+
   // Finally put everything together.
   os << llvm::formatv(header, cppOpName, linalgOpName, interfaceNameList, doc,
-                      attrList, state.orderedTensorArgs.size(), attrBuilder);
+                      attrList, state.orderedTensorArgs.size(), attrBuilder,
+                      attrMethods);
 }
 
 /// Print the C++ StructuredOpsInterface impl of `iterator_types`.
 void TCParser::printReferenceIterators(llvm::raw_ostream &os,
                                        StringRef cppOpName,
                                        ComprehensionParsingState &state) {
   const char *referenceReferenceIteratorsFmt =
       R"FMT(
const char *canonicalizersAndFoldersFmt = R"FMT(
     void {0}::getEffects(SmallVectorImpl<
         SideEffects::EffectInstance<MemoryEffects::Effect> >&effects) {{
       getGenericEffectsImpl(effects,
         getOperation()->getResults(), getInputBuffers(), getOutputBuffers());
     })FMT";
   os << llvm::formatv(canonicalizersAndFoldersFmt, cppOpName);
 }
 
+// Prints methods for querying whether the current named op has attributes that
+// are used by its indexing maps and for verifying those attributes have the
+// expected type.
+void TCParser::printIndexingMapRequiredAttrMethods(
+    llvm::raw_ostream &os, StringRef cppOpName,
+    ComprehensionParsingState &state) {
+  // If there are no attribute used by the whole definition, then we are done.
+  if (registeredAttrs.empty())
+    return;
+
+  // Otherwise, go through each attribute and generate code to verify it's
+  // valid per the spec.
+  SmallVector<std::string, 4> attributes;
+  for (const auto &attr : registeredAttrs) {
+    if (attr.second.isOptional)
+      continue;
+
+    llvm::StringRef name = attr.first;
+    llvm::StringRef elementType = attr.second.elementType;
+    const auto &dims = attr.second.vectorDims;
+
+    // Get the method call to check the element type is of the expected kind.
+    std::string elemTypeCheck = llvm::StringSwitch<std::string>(elementType)
+                                    .Case("f32", "isF32()")
+                                    .Case("i32", "isInteger(32)")
+                                    .Case("i64", "isInteger(64)")
+                                    .Default("");
+    if (elemTypeCheck.empty()) {
+      (void)parser.emitError(
+          "unimplemented support for attribute element type: " + elementType);
+      return;
+    }
+
+    // Scalar case.
+    if (dims.empty() && !attr.second.isArray) {
+      const char *attrFmt = R"FMT(
+        if (auto attr = op->getAttr("{0}")) {{
+          if (!attr.getType().{1}) return op->emitError(
+            "incorrect type for indexing map required attribute '{0}'");
+        } else {{
+          return op->emitError(
+            "missing indexing map required attribute '{0}'");
+        }
+      )FMT";
+
+      attributes.push_back(llvm::formatv(attrFmt, name, elemTypeCheck));
+      continue;
+    }
+
+    // Vector case.
+    if (!dims.empty()) {
+      SmallVector<std::string, 4> dimStrs;
+      for (uint64_t dim : dims)
+        dimStrs.push_back(std::to_string(dim));
+
+      const char *attrFmt = R"FMT(
+        if (auto attr = op->getAttrOfType<DenseElementsAttr>("{0}")) {{
+          if (!attr.getType().getElementType().{1}) return op->emitError(
+            "incorrect element type for indexing map required attribute '{0}'");
+          if (attr.getType().getShape() != ArrayRef<int64_t>{{ {2} })
+            return op->emitError(
+              "incorrect shape for indexing map required attribute '{0}'");
+        } else {
+          return op->emitError(
+            "missing indexing map required attribute '{0}'");
+        }
+      )FMT";
+
+      attributes.push_back(llvm::formatv(attrFmt, name, elemTypeCheck,
+                                         llvm::join(dimStrs, ", ")));
+      continue;
+    }
+
+    // Array case.
+    {
+      const char *attrFmt = R"FMT(
+        if (auto attr = op->getAttrOfType<ArrayAttr>("{0}")) {{
+          for (Attribute element : attr) {{
+            if (!element.getType().{1}) return emitError(
+              "incorrect element type for indexing map required attribute '{0}'");
+          }
+        } else {{
+          return op->emitError(
+            "missing indexing map required attribute '{0}'");
+        }
+      )FMT";
+
+      attributes.push_back(llvm::formatv(attrFmt, name, elemTypeCheck));
+    }
+  }
+
+  const char *methodFmt = R"FMT(
+  bool {0}::hasDynamicIndexingMaps() {{ return true; }
+
+  LogicalResult {0}::verifyIndexingMapRequiredAttributes() {{
+    Operation *op = getOperation();
+    {1}
+    return success();
+  }
+  )FMT";
+
+  // Print everything out.
+  os << llvm::formatv(methodFmt, cppOpName, llvm::join(attributes, "\n"));
+}
+
 /// Print the C++ StructuredOpsInterface impl of `referenceIndexingMaps`.
 void TCParser::printReferenceIndexingMaps(llvm::raw_ostream &os,
                                           StringRef cppOpName,
                                           ComprehensionParsingState &state) {
   // 1. Generic string template for specifying reference indexing maps.
   const char *referenceIndexingMapsFmt =
       R"FMT(
   // This is temporary until we transition out of manually specified ops that