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

   1 //===- mlir-linalg-ods-yaml-gen.cpp - Linalg ODS generation from yaml  ----===//
   2 //
   3 // Part of the LLVM Project, 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 // This file implements an ODS (and C++) generator from a YAML form
  10 // derived from the mathematical expression of linalg named ops. Typically a
  11 // math oriented DSL will be used to export the essential representation to
  12 // this form, and maintaining the SOT at the math level (versus recreating it
  13 // in MLIR) is deemed to have systemic value.
  14 //
  15 //===----------------------------------------------------------------------===//
  16
  17 #include "mlir/AsmParser/AsmParser.h"
  18 #include "mlir/IR/AffineMap.h"
  19 #include "mlir/IR/Diagnostics.h"
  20 #include "mlir/IR/MLIRContext.h"
  21 #include "mlir/Support/FileUtilities.h"
  22 #include "mlir/Support/LLVM.h"
  23 #include "llvm/ADT/StringRef.h"
  24 #include "llvm/Support/CommandLine.h"
  25 #include "llvm/Support/Debug.h"
  26 #include "llvm/Support/FormatVariadic.h"
  27 #include "llvm/Support/ToolOutputFile.h"
  28 #include "llvm/Support/YAMLTraits.h"
  29 #include <optional>
  30
  31 using namespace mlir;
  32
  33 using llvm::yaml::Input;
  34 using llvm::yaml::MappingTraits;
  35 using llvm::yaml::ScalarEnumerationTraits;
  36 using llvm::yaml::ScalarTraits;
  37
  38 #define DEBUG_TYPE "linalg-ods-gen"
  39
  40 //===----------------------------------------------------------------------===//
  41 // Mapping structs (correspond to data types in the YAML description).
  42 // TODO: Since this is a schema/part of the contract, it should be moved to
  43 // a real header.
  44 //===----------------------------------------------------------------------===//
  45
  46 namespace {
  47
  48 struct LinalgYAMLContext {
  49   MLIRContext *mlirContext;
  50 };
  51
  52 struct LinalgOpMetadata {
  53   std::string name;
  54   std::string cppClassName;
  55   std::optional<std::string> doc;
  56   SmallVector<std::string> implements;
  57   SmallVector<std::string> defines;
  58 };
  59
  60 struct SerializedAffineMap {
  61   AffineMapAttr affineMapAttr;
  62
  63   AffineMap affineMap() { return affineMapAttr.getValue(); }
  64 };
  65
  66 enum class LinalgOperandDefKind {
  67   InputTensor,
  68   Scalar,
  69   OutputTensor,
  70   IndexAttr,
  71   UnaryFnAttr,
  72   BinaryFnAttr,
  73   TypeFnAttr
  74 };
  75
  76 struct LinalgOperandDef {
  77   std::string name;
  78   LinalgOperandDefKind kind;
  79   std::optional<std::string> typeVar;
  80   std::optional<SerializedAffineMap> shapeMap;
  81   std::optional<SerializedAffineMap> indexAttrMap;
  82   std::optional<SmallVector<int64_t>> defaultIndices;
  83   std::optional<std::string> defaultFn;
  84 };
  85
  86 enum class LinalgIteratorTypeDef {
  87   parallel,
  88   reduction,
  89 };
  90
  91 struct LinalgIndexingMapsConfig {
  92   std::optional<SmallVector<SerializedAffineMap>> staticIndexingMaps;
  93 };
  94
  95 struct ScalarExpression;
  96
  97 enum class ScalarFnKind { Unary, Binary, Type };
  98
  99 struct ScalarFn {
 100   ScalarFnKind kind;
 101   std::optional<std::string> fnName;
 102   std::optional<std::string> attrName;
 103   std::optional<std::string> typeVar;
 104   // NOTE: This must be of arity 1, but to break the self-referential cycle,
 105   // we use a heap allocated vector.
 106   std::vector<ScalarExpression> operands;
 107 };
 108
 109 struct ScalarExpression {
 110   std::optional<std::string> arg;
 111   std::optional<std::string> constant;
 112   std::optional<int64_t> index;
 113   std::optional<ScalarFn> scalarFn;
 114 };
 115
 116 struct ScalarAssign {
 117   std::string arg;
 118   ScalarExpression value;
 119 };
 120
 121 struct LinalgStructuredOpConfig {
 122   SmallVector<LinalgOperandDef> args;
 123   LinalgIndexingMapsConfig indexingMaps;
 124   SmallVector<LinalgIteratorTypeDef> iteratorTypes;
 125   std::vector<ScalarAssign> assignments;
 126 };
 127
 128 struct LinalgOpConfig {
 129   std::optional<LinalgOpMetadata> metadata;
 130   std::optional<LinalgStructuredOpConfig> structuredOp;
 131 };
 132
 133 } // namespace
 134
 135 //===----------------------------------------------------------------------===//
 136 // Mapping traits.
 137 //===----------------------------------------------------------------------===//
 138
 139 LLVM_YAML_IS_SEQUENCE_VECTOR(LinalgOperandDef)
 140 LLVM_YAML_IS_SEQUENCE_VECTOR(SerializedAffineMap)
 141 LLVM_YAML_IS_SEQUENCE_VECTOR(LinalgIteratorTypeDef)
 142 LLVM_YAML_IS_SEQUENCE_VECTOR(ScalarAssign)
 143 LLVM_YAML_IS_SEQUENCE_VECTOR(ScalarExpression)
 144 LLVM_YAML_IS_DOCUMENT_LIST_VECTOR(LinalgOpConfig)
 145
 146 namespace llvm {
 147 namespace yaml {
 148
 149 /// Top-level type containing op metadata and one of a concrete op type.
 150 /// Currently, the only defined op type is `structured_op` (maps to
 151 /// `LinalgStructuredOpConfig`).
 152 template <>
 153 struct MappingTraits<LinalgOpConfig> {
 154   static void mapping(IO &io, LinalgOpConfig &info) {
 155     io.mapOptional("metadata", info.metadata);
 156     io.mapOptional("structured_op", info.structuredOp);
 157   }
 158 };
 159
 160 /// A structured op models (at most) a single contraction by modeling
 161 ///   - A list of named arguments (`LinalgOperandDef`), which can be inputs,
 162 ///     outputs, or index attributes.
 163 ///   - List of indexing maps (see `LinalgIndexingMaps`).
 164 ///   - Iterator types (see `LinalgIteratorTypeDef`).
 165 ///   - List of scalar level assignment (see `ScalarAssign`).
 166 template <>
 167 struct MappingTraits<LinalgStructuredOpConfig> {
 168   static void mapping(IO &io, LinalgStructuredOpConfig &info) {
 169     io.mapRequired("args", info.args);
 170     io.mapRequired("indexing_maps", info.indexingMaps);
 171     io.mapRequired("iterator_types", info.iteratorTypes);
 172     io.mapRequired("assignments", info.assignments);
 173   }
 174 };
 175
 176 /// Maps a named tensor, scalar or attribute argument to an operation,
 177 /// consisting of:
 178 ///   - `name`: Must be unique within the operation.
 179 ///   - `usage`: How the argument is used (input, output, attribute, etc).
 180 ///   - `type_var`: The symbolic type variable that binds to the element or self
 181 ///     type of the tensor or scalar argument, respectively.
 182 ///   - `shape_map`: An optional AffineMap from all op symbols to the shape of
 183 ///     the argument. Only tensor arguments have a `shape_map`. Each shape must
 184 ///     be normalized over the same list of symbols and have no dimension
 185 ///     inputs.
 186 ///   - `index_attr_map`: An optional AffineMap from all op symbols to the
 187 ///     index attribute symbols. During op creation these symbols are replaced
 188 ///     by the corresponding `name` index attribue values. Only index attribute
 189 ///     arguments have an `index_attr_map`.
 190 ///   - `default_indices`: An optional default initialization for index
 191 ///     attribute arguments.
 192 ///   - `default_fn`: An optional default initialization for function attribute
 193 ///     arguments.
 194 template <>
 195 struct MappingTraits<LinalgOperandDef> {
 196   static void mapping(IO &io, LinalgOperandDef &info) {
 197     io.mapRequired("name", info.name);
 198     io.mapRequired("kind", info.kind);
 199     io.mapOptional("type_var", info.typeVar);
 200     io.mapOptional("shape_map", info.shapeMap);
 201     io.mapOptional("index_attr_map", info.indexAttrMap);
 202     io.mapOptional("default_indices", info.defaultIndices);
 203     io.mapOptional("default_fn", info.defaultFn);
 204   }
 205 };
 206
 207 /// Usage enum for a named argument.
 208 template <>
 209 struct ScalarEnumerationTraits<LinalgOperandDefKind> {
 210   static void enumeration(IO &io, LinalgOperandDefKind &value) {
 211     io.enumCase(value, "input_tensor", LinalgOperandDefKind::InputTensor);
 212     io.enumCase(value, "scalar", LinalgOperandDefKind::Scalar);
 213     io.enumCase(value, "output_tensor", LinalgOperandDefKind::OutputTensor);
 214     io.enumCase(value, "index_attr", LinalgOperandDefKind::IndexAttr);
 215     io.enumCase(value, "unary_fn_attr", LinalgOperandDefKind::UnaryFnAttr);
 216     io.enumCase(value, "binary_fn_attr", LinalgOperandDefKind::BinaryFnAttr);
 217     io.enumCase(value, "type_fn_attr", LinalgOperandDefKind::TypeFnAttr);
 218   }
 219 };
 220
 221 /// Iterator type enum.
 222 template <>
 223 struct ScalarEnumerationTraits<LinalgIteratorTypeDef> {
 224   static void enumeration(IO &io, LinalgIteratorTypeDef &value) {
 225     io.enumCase(value, "parallel", LinalgIteratorTypeDef::parallel);
 226     io.enumCase(value, "reduction", LinalgIteratorTypeDef::reduction);
 227   }
 228 };
 229
 230 /// Metadata about the op (name, C++ name, and documentation).
 231 template <>
 232 struct MappingTraits<LinalgOpMetadata> {
 233   static void mapping(IO &io, LinalgOpMetadata &info) {
 234     io.mapRequired("name", info.name);
 235     io.mapRequired("cpp_class_name", info.cppClassName);
 236     io.mapOptional("doc", info.doc);
 237     io.mapOptional("implements", info.implements);
 238     io.mapOptional("defines", info.defines);
 239   }
 240 };
 241
 242 /// How the ops indexing maps are produced. Must be one of:
 243 ///   - static_indexing_maps: A static list of AffineMaps, possibly with
 244 ///     some symbols that bind to attributes of the op. Each indexing map must
 245 ///     be normalized over the same list of dimensions, and its symbols must
 246 ///     match the symbols for argument shapes.
 247 template <>
 248 struct MappingTraits<LinalgIndexingMapsConfig> {
 249   static void mapping(IO &io, LinalgIndexingMapsConfig &info) {
 250     io.mapOptional("static_indexing_maps", info.staticIndexingMaps);
 251   }
 252 };
 253
 254 /// Models an assignment to a named output.
 255 ///   - The `arg` name must match a named output.
 256 ///   - The `value` is a scalar expression for computing the value to
 257 ///     assign (see `ScalarExpression`).
 258 template <>
 259 struct MappingTraits<ScalarAssign> {
 260   static void mapping(IO &io, ScalarAssign &info) {
 261     io.mapRequired("arg", info.arg);
 262     io.mapRequired("value", info.value);
 263   }
 264 };
 265
 266 /// A scalar expression (RHS of an assignment). Must be one of:
 267 ///   - `scalar_arg`: An operation argument.
 268 ///   - `scalar_const`: A constant definition.
 269 ///   - `scalar_index`: An iteration index.
 270 ///   - `scalar_fn`: A named function (see `ScalarFn`).
 271 template <>
 272 struct MappingTraits<ScalarExpression> {
 273   static void mapping(IO &io, ScalarExpression &info) {
 274     io.mapOptional("scalar_arg", info.arg);
 275     io.mapOptional("scalar_const", info.constant);
 276     io.mapOptional("scalar_index", info.index);
 277     io.mapOptional("scalar_fn", info.scalarFn);
 278   }
 279 };
 280
 281 /// Scalar function kind enum.
 282 template <>
 283 struct ScalarEnumerationTraits<ScalarFnKind> {
 284   static void enumeration(IO &io, ScalarFnKind &value) {
 285     io.enumCase(value, "unary", ScalarFnKind::Unary);
 286     io.enumCase(value, "binary", ScalarFnKind::Binary);
 287     io.enumCase(value, "type", ScalarFnKind::Type);
 288   }
 289 };
 290
 291 /// A scalar expression that evaluates a named function.
 292 /// Functions are generally "math" level and type polymorphic. Builtin
 293 /// functions include:
 294 ///   - `add(lhs, rhs)`
 295 ///   - `mul(lhs, rhs)`
 296 template <>
 297 struct MappingTraits<ScalarFn> {
 298   static void mapping(IO &io, ScalarFn &info) {
 299     io.mapRequired("kind", info.kind);
 300     io.mapOptional("fn_name", info.fnName);
 301     io.mapOptional("attr_name", info.attrName);
 302     io.mapOptional("type_var", info.typeVar);
 303     io.mapRequired("operands", info.operands);
 304   }
 305 };
 306
 307 /// Helper mapping which accesses an AffineMapAttr as a serialized string of
 308 /// the same.
 309 template <>
 310 struct ScalarTraits<SerializedAffineMap> {
 311   static void output(const SerializedAffineMap &value, void *rawYamlContext,
 312                      raw_ostream &out) {
 313     assert(value.affineMapAttr);
 314     value.affineMapAttr.print(out);
 315   }
 316   static StringRef input(StringRef scalar, void *rawYamlContext,
 317                          SerializedAffineMap &value) {
 318     assert(rawYamlContext);
 319     auto *yamlContext = static_cast<LinalgYAMLContext *>(rawYamlContext);
 320     if (auto attr = dyn_cast_or_null<AffineMapAttr>(
 321             mlir::parseAttribute(scalar, yamlContext->mlirContext)))
 322       value.affineMapAttr = attr;
 323     else if (!value.affineMapAttr || !isa<AffineMapAttr>(value.affineMapAttr))
 324       return "could not parse as an affine map attribute";
 325     return StringRef();
 326   }
 327   static QuotingType mustQuote(StringRef) { return QuotingType::None; }
 328 };
 329
 330 } // namespace yaml
 331 } // namespace llvm
 332
 333 namespace {
 334
 335 //===----------------------------------------------------------------------===//
 336 // Generation utilities
 337 //===----------------------------------------------------------------------===//
 338
 339 class GenerationContext {
 340 public:
 341   GenerationContext(MLIRContext *context, raw_ostream *odsOut,
 342                     raw_ostream *defnOut)
 343       : context(context), loc(UnknownLoc::get(context)), odsOut(odsOut),
 344         defnOut(defnOut) {}
 345
 346   MLIRContext *getContext() { return context; }
 347
 348   void setLoc(Location loc) { this->loc = loc; }
 349   Location getLoc() { return loc; }
 350
 351   bool shouldGenerateOds() { return odsOut; }
 352   bool shouldGenerateDefns() { return defnOut; }
 353
 354   raw_ostream &odss() {
 355     assert(odsOut && "ODS stream not defined");
 356     return *odsOut;
 357   }
 358
 359   raw_ostream &defns() {
 360     assert(defnOut && "Definition stream not defined");
 361     return *defnOut;
 362   }
 363
 364 private:
 365   MLIRContext *context;
 366   Location loc;
 367   raw_ostream *odsOut;
 368   raw_ostream *defnOut;
 369 };
 370
 371 } // namespace
 372
 373 static std::string generateCppExpression(SerializedAffineMap self,
 374                                          StringRef contextName) {
 375   std::string printedStr;
 376   llvm::raw_string_ostream printedSs(printedStr);
 377   self.affineMapAttr.print(printedSs);
 378   printedSs.flush();
 379
 380   static const char exprFormat[] =
 381       R"FMT(llvm::cast<AffineMapAttr>(mlir::parseAttribute("{0}", {1})).getValue())FMT";
 382   return llvm::formatv(exprFormat, printedStr, contextName);
 383 }
 384
 385 template <typename Container>
 386 static std::string interleaveToString(Container &container,
 387                                       StringRef separator) {
 388   std::string result;
 389   llvm::raw_string_ostream ss(result);
 390   llvm::interleave(container, ss, separator);
 391   ss.flush();
 392   return result;
 393 }
 394
 395 static std::optional<int>
 396 findTensorDefArgIndex(StringRef name, SmallVectorImpl<LinalgOperandDef> &args) {
 397   for (const auto &it : llvm::enumerate(args)) {
 398     if (it.value().name == name)
 399       return it.index();
 400   }
 401   return std::nullopt;
 402 }
 403
 404 // Try to map the TypeVar to a predefined or an argument type.
 405 static std::optional<std::string>
 406 findTypeValue(StringRef typeVar, SmallVectorImpl<LinalgOperandDef> &args) {
 407   // Handle all predefined types.
 408   if (typeVar == "I32")
 409     return std::string("helper.getIntegerType(32)");
 410   if (typeVar == "I64")
 411     return std::string("helper.getIntegerType(64)");
 412   if (typeVar == "F32")
 413     return std::string("helper.getFloat32Type()");
 414   if (typeVar == "F64")
 415     return std::string("helper.getFloat64Type()");
 416
 417   // Search all argument types.
 418   for (const auto &it : llvm::enumerate(args)) {
 419     if (it.value().kind != LinalgOperandDefKind::InputTensor &&
 420         it.value().kind != LinalgOperandDefKind::Scalar &&
 421         it.value().kind != LinalgOperandDefKind::OutputTensor)
 422       continue;
 423     if (*it.value().typeVar == typeVar)
 424       return llvm::formatv("block.getArgument({0}).getType()", it.index())
 425           .str();
 426   }
 427
 428   return std::nullopt;
 429 }
 430
 431 static ScalarAssign *findAssignment(StringRef name,
 432                                     std::vector<ScalarAssign> &assignments) {
 433   for (auto &assign : assignments) {
 434     if (assign.arg == name)
 435       return &assign;
 436   }
 437   return nullptr;
 438 }
 439
 440 // Return true if the operand is a function attribute.
 441 static bool isFunctionAttribute(LinalgOperandDefKind kind) {
 442   return kind == LinalgOperandDefKind::UnaryFnAttr ||
 443          kind == LinalgOperandDefKind::BinaryFnAttr ||
 444          kind == LinalgOperandDefKind::TypeFnAttr;
 445 }
 446
 447 // Return true if the operand is an attribute.
 448 static bool isAttribute(LinalgOperandDefKind kind) {
 449   return kind == LinalgOperandDefKind::IndexAttr || isFunctionAttribute(kind);
 450 }
 451
 452 // Get the enum name for the given operand kind.
 453 std::string convertOperandKindToEnumName(LinalgOperandDefKind kind) {
 454   switch (kind) {
 455   case LinalgOperandDefKind::UnaryFnAttr:
 456     return std::string("UnaryFn");
 457   case LinalgOperandDefKind::BinaryFnAttr:
 458     return std::string("BinaryFn");
 459   case LinalgOperandDefKind::TypeFnAttr:
 460     return std::string("TypeFn");
 461   default:
 462     break;
 463   }
 464   llvm_unreachable("unsupported function attribute kind");
 465 }
 466
 467 // Get the enum name for the given function kind.
 468 std::string convertFunctionKindToEnumName(ScalarFnKind kind) {
 469   switch (kind) {
 470   case ScalarFnKind::Unary:
 471     return std::string("UnaryFn");
 472   case ScalarFnKind::Binary:
 473     return std::string("BinaryFn");
 474   case ScalarFnKind::Type:
 475     return std::string("TypeFn");
 476   }
 477   llvm_unreachable("unsupported function kind");
 478 }
 479
 480 //===----------------------------------------------------------------------===//
 481 // Templates
 482 //===----------------------------------------------------------------------===//
 483
 484 // A single line banner format. Parameters:
 485 // {0}: Single line comment
 486 static const char bannerFormat[] = R"FMT(
 487 //===----------------------------------------------------------------------===//
 488 // {0}
 489 //===----------------------------------------------------------------------===//
 490 )FMT";
 491
 492 //===----------------------------------------------------------------------===//
 493 // Named generic op generation.
 494 // These ops map at most a single contraction that complies with the limitations
 495 // of a linalg.generic.
 496 //===----------------------------------------------------------------------===//
 497
 498 // Template for Linalg named ops' ODS definitions. Parameters:
 499 // {0}: ODS/C++ op name
 500 // {1}: assembly op mnemonic
 501 // {2}: op interface list
 502 // {3}: documentation (summary + description)
 503 // {4}: op attribute list
 504 // {5}: builder methods taking standalone attribute parameters
 505 // {6}: additional method defintions
 506 // {7}: additional methods for attributes used by indexing maps
 507 static const char structuredOpOdsHeaderFormat[] = R"FMT(
 508 //===----------------------------------------------------------------------===//
 509 // Op definition for {0}
 510 //===----------------------------------------------------------------------===//
 511
 512 def {0} : LinalgStructuredBase_Op<"{1}", !listconcat([AttrSizedOperandSegments],
 513   /*extraInterfaces=*/[{2}])> {
 514     {3}
 515     let arguments = (ins
 516       Variadic<AnyType>:$inputs,
 517       Variadic<AnyShaped>:$outputs{4}
 518     );
 519     let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
 520     let regions = (region AnyRegion:$region);
 521
 522     let skipDefaultBuilders = 1;
 523     let builders = [
 524       OpBuilder<
 525       (ins "ValueRange":$inputs, "ValueRange":$outputs,
 526             CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
 527       [{{
 528         buildStructuredOp($_builder, $_state, std::nullopt, inputs, outputs,
 529           attributes, {0}::getRegionBuilder());
 530       }]>,
 531       OpBuilder<
 532       (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
 533             "ValueRange":$outputs,
 534             CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
 535       [{{
 536         buildStructuredOp($_builder, $_state, resultTensorTypes,
 537           inputs, outputs, attributes, {0}::getRegionBuilder());
 538       }]>,
 539       OpBuilder<
 540       (ins "TypeRange":$resultTensorTypes, "ValueRange":$operands,
 541             CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
 542       [{{
 543         $_state.addOperands(operands);
 544         $_state.addAttributes(attributes);
 545         $_state.addTypes(resultTensorTypes);
 546         (void)$_state.addRegion();
 547       }]>
 548       {5}
 549     ];
 550     let hasCustomAssemblyFormat = 1;
 551     let hasFolder = 1;
 552     {6}
 553
 554     let extraClassDeclaration = structuredOpsBaseDecls # [{{
 555       // Auto-generated.
 556       SmallVector<utils::IteratorType> getIteratorTypesArray();
 557       ArrayAttr getIndexingMaps();
 558       static void regionBuilder(ImplicitLocOpBuilder &b,
 559                                 Block &block, ArrayRef<NamedAttribute> attrs);
 560       static std::function<void(ImplicitLocOpBuilder &,
 561                                 Block &, ArrayRef<NamedAttribute>)>
 562       getRegionBuilder() {{
 563         return regionBuilder;
 564       }
 565
 566       std::pair<int64_t, int64_t> getDpsInitsPositionRange() {{
 567         int64_t getNumOperands = this->getNumOperands();
 568         return {{getNumOperands - getOutputs().size(), getNumOperands};
 569       }
 570
 571       // Generic methods.
 572       static unsigned getNumRegionArgs();
 573       std::string getLibraryCallName();
 574       {7}
 575     }];
 576 }
 577 )FMT";
 578
 579 // Builder method taking attribute parameters. Parameters:
 580 // {0}: Class name
 581 // {1}: Comma interleaved attribute parameters
 582 // {2}: Attribute initialization
 583 static const char structuredOpBuilderFormat[] = R"FMT(
 584   , OpBuilder<
 585   (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
 586        "ValueRange":$outputs, {1},
 587        CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
 588   [{{
 589     {2}
 590     buildStructuredOp($_builder, $_state, resultTensorTypes, inputs, outputs,
 591       attributes, {0}::getRegionBuilder());
 592   }]>
 593 )FMT";
 594
 595 // The getIteratorTypesArray() method for structured ops. Parameters:
 596 // {0}: Class name
 597 // {1}: Comma interleaved iterator type names.
 598 static const char structuredOpIteratorTypesFormat[] =
 599     R"FMT(
 600 SmallVector<utils::IteratorType> {0}::getIteratorTypesArray() {{
 601   return SmallVector<utils::IteratorType>{{ {1} };
 602 }
 603 )FMT";
 604
 605 // The getIteratorTypesArray() method for rank polymorphic structured ops.
 606 // Parameters:
 607 // {0}: Class name
 608 static const char rankPolyStructuredOpIteratorTypesFormat[] =
 609     R"FMT(
 610 SmallVector<utils::IteratorType> {0}::getIteratorTypesArray() {{
 611   int64_t rank = getRank(getDpsInitOperand(0));
 612   return SmallVector<utils::IteratorType>(rank, utils::IteratorType::parallel);
 613 }
 614 )FMT";
 615
 616 // The indexing_maps() method for structured ops. Parameters:
 617 // {0}: Class name
 618 // {1}: Comma-separated list of dimension variable names.
 619 // {2}: Statements
 620 static const char structuredOpIndexingMapsFormat[] = R"FMT(
 621 ArrayAttr {0}::getIndexingMaps() {{
 622   static const char memoizeAttr[] = "linalg.memoized_indexing_maps";
 623   ArrayAttr cached = getOperation()->getAttrOfType<ArrayAttr>(memoizeAttr);
 624   if (cached)
 625     return cached;
 626
 627   MLIRContext *context = getContext();
 628   auto symbolBindings = getSymbolBindings(*this);
 629   SmallVector<AffineMap> maps;
 630   {2}
 631   cached = Builder(context).getAffineMapArrayAttr(maps);
 632   getOperation()->setAttr(memoizeAttr, cached);
 633   return cached;
 634 }
 635 )FMT";
 636
 637 // The indexing_maps() method for rank polymorphic structured ops. Parameters:
 638 // {0}: Class name
 639 static const char rankPolyStructuredOpIndexingMapsFormat[] = R"FMT(
 640 ArrayAttr {0}::getIndexingMaps() {{
 641   MLIRContext *context = getContext();
 642   AffineMap scalarMap = AffineMap::get(getNumParallelLoops(), 0, context);
 643   AffineMap tensorMap = AffineMap::getMultiDimIdentityMap(
 644     getNumParallelLoops(), context);
 645   SmallVector<AffineMap> indexingMaps;
 646   for (OpOperand &opOperand : getOperation()->getOpOperands())
 647     indexingMaps.push_back(getRank(&opOperand) == 0 ? scalarMap : tensorMap);
 648   return Builder(getContext()).getAffineMapArrayAttr(indexingMaps);
 649 }
 650 )FMT";
 651
 652 // Implementations of fold and getEffects.
 653 // Parameters:
 654 // {0}: Class name
 655 const char structuredOpFoldersFormat[] = R"FMT(
 656 LogicalResult {0}::fold(FoldAdaptor,
 657                         SmallVectorImpl<OpFoldResult> &) {{
 658   return memref::foldMemRefCast(*this);
 659 }
 660 void {0}::getEffects(SmallVectorImpl<
 661     SideEffects::EffectInstance<MemoryEffects::Effect> >&effects) {{
 662       if (hasTensorSemantics()) return;
 663       getGenericEffectsImpl(effects,
 664         getOperation()->getResults(), getDpsInputOperands(), getDpsInitOperands());
 665 }
 666 )FMT";
 667
 668 // Implementation of parse/print.
 669 // Parameters:
 670 // {0}: Class name
 671 static const char structuredOpParserFormat[] = R"FMT(
 672 ParseResult {0}::parse(OpAsmParser &parser, OperationState &result) {{
 673   return ::parseNamedStructuredOp(parser, result,
 674     {0}::getNumRegionArgs(), {0}::getRegionBuilder());
 675 }
 676 void {0}::print(OpAsmPrinter &p) {{
 677   ::printNamedStructuredOp(p, getOperation(), getInputs(), getOutputs());
 678 }
 679 )FMT";
 680
 681 static LogicalResult generateNamedGenericOpOds(LinalgOpConfig &opConfig,
 682                                                GenerationContext &genContext) {
 683   if (!genContext.shouldGenerateOds())
 684     return success();
 685
 686   raw_ostream &os = genContext.odss();
 687
 688   std::string interfaceNameList;
 689   std::string attrList;
 690   std::string attrMethods;
 691   std::string attrBuilder;
 692
 693   std::string doc;
 694   if (opConfig.metadata->doc) {
 695     static const char structuredOpDocFmt[] = R"FMT(
 696   let summary = [{{{0}}];
 697   let description = [{{{1}}];
 698 )FMT";
 699     StringRef summary, description;
 700     std::tie(summary, description) =
 701         StringRef(*opConfig.metadata->doc).trim().split("\n\n");
 702
 703     doc = llvm::formatv(structuredOpDocFmt, summary.trim(), description.trim());
 704   }
 705
 706   interfaceNameList = interleaveToString(opConfig.metadata->implements, ", ");
 707
 708   std::string definitionList;
 709   for (const std::string &definition : opConfig.metadata->defines) {
 710     static const char definitionFmt[] = "let {0} = 1;\n";
 711     definitionList.append(llvm::formatv(definitionFmt, definition));
 712   }
 713
 714   if (llvm::any_of(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
 715         return isAttribute(arg.kind);
 716       })) {
 717     SmallVector<std::string> attrDefs;
 718     SmallVector<std::string> attrParams;
 719     SmallVector<std::string> attrStmts;
 720     for (LinalgOperandDef &arg : opConfig.structuredOp->args) {
 721       static const char paramFmt[] = "\"Attribute\":${0}";
 722       static const char stmtFmt[] = "$_state.addAttribute(\"{0}\", {0});";
 723       // Add the type conversion attributes to the op definition and builders.
 724       if (isFunctionAttribute(arg.kind)) {
 725         assert(arg.defaultFn);
 726         std::string enumName = convertOperandKindToEnumName(arg.kind);
 727         static const char typeFmt[] = "{0}::{1}";
 728         static const char defFmt[] =
 729             "DefaultValuedOptionalAttr<{0}, \"{1}\">:${2}";
 730         attrDefs.push_back(llvm::formatv(
 731             defFmt, llvm::formatv("{0}Attr", enumName),
 732             llvm::formatv(typeFmt, enumName, arg.defaultFn), arg.name));
 733         attrParams.push_back(llvm::formatv(paramFmt, arg.name));
 734         attrStmts.push_back(llvm::formatv(stmtFmt, arg.name));
 735       }
 736       // Add the index attributes to the op definition and builders.
 737       if (arg.kind == LinalgOperandDefKind::IndexAttr) {
 738         assert(arg.indexAttrMap.has_value());
 739         assert(arg.defaultIndices.has_value());
 740         size_t size = arg.indexAttrMap->affineMap().getNumResults();
 741         assert(arg.defaultIndices->size() == size);
 742         static const char typeFmt[] = "RankedI64ElementsAttr<[{0}]>";
 743         static const char defFmt[] =
 744             "DefaultValuedOptionalAttr<{0}, \"{ {1} }\">:${2}";
 745         std::string defaultVals;
 746         llvm::raw_string_ostream ss(defaultVals);
 747         llvm::interleave(
 748             *arg.defaultIndices, ss,
 749             [&](int64_t val) { ss << "static_cast<int64_t>(" << val << ")"; },
 750             ", ");
 751         attrDefs.push_back(llvm::formatv(defFmt, llvm::formatv(typeFmt, size),
 752                                          ss.str(), arg.name));
 753         attrParams.push_back(llvm::formatv(paramFmt, arg.name));
 754         attrStmts.push_back(llvm::formatv(stmtFmt, arg.name));
 755       }
 756     }
 757     if (llvm::any_of(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
 758           return arg.kind == LinalgOperandDefKind::IndexAttr;
 759         })) {
 760       attrMethods = R"(
 761         bool hasDynamicIndexingMaps();
 762         LogicalResult verifyIndexingMapRequiredAttributes();
 763       )";
 764     }
 765     attrList = ",\n" + llvm::join(attrDefs, ",\n");
 766     attrBuilder = llvm::formatv(
 767         structuredOpBuilderFormat, opConfig.metadata->cppClassName,
 768         llvm::join(attrParams, ", "), llvm::join(attrStmts, "\n"));
 769   }
 770
 771   os << llvm::formatv(structuredOpOdsHeaderFormat,
 772                       opConfig.metadata->cppClassName, opConfig.metadata->name,
 773                       interfaceNameList, doc, attrList, attrBuilder,
 774                       definitionList, attrMethods);
 775
 776   return success();
 777 }
 778
 779 static LogicalResult
 780 generateNamedGenericOpDefns(LinalgOpConfig &opConfig,
 781                             GenerationContext &genContext) {
 782   if (!genContext.shouldGenerateDefns())
 783     return success();
 784
 785   raw_ostream &os = genContext.defns();
 786   StringRef className = opConfig.metadata->cppClassName;
 787
 788   // Implementation banner.
 789   std::string bannerComment = llvm::formatv("Implementation of {0}", className);
 790   os << llvm::formatv(bannerFormat, bannerComment);
 791
 792   // Compute the number of scalar and tensor arguments.
 793   int64_t numOfArgs =
 794       llvm::count_if(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
 795         return arg.kind == LinalgOperandDefKind::InputTensor ||
 796                arg.kind == LinalgOperandDefKind::Scalar ||
 797                arg.kind == LinalgOperandDefKind::OutputTensor;
 798       });
 799
 800   // An operation that accesses only scalars and scalar/rank zero tensors is
 801   // rank polymorhpic. We implement rank polymorphism by generating different
 802   // indexing maps and iterators that match the rank of the first output tensor.
 803   // An operation is rank polymorphic if the iteration domain has rank zero.
 804   bool isRankPolymorphic = opConfig.structuredOp->iteratorTypes.empty();
 805
 806   // Generate the iterator_types() method.
 807   if (!isRankPolymorphic) {
 808     std::string iteratorsStr;
 809     llvm::raw_string_ostream ss(iteratorsStr);
 810     llvm::interleaveComma(opConfig.structuredOp->iteratorTypes, ss,
 811                           [&](LinalgIteratorTypeDef it) {
 812                             switch (it) {
 813                             case LinalgIteratorTypeDef::parallel:
 814                               ss << "utils::IteratorType::parallel";
 815                               break;
 816                             case LinalgIteratorTypeDef::reduction:
 817                               ss << "utils::IteratorType::reduction";
 818                               break;
 819                             }
 820                           });
 821     ss.flush();
 822     os << llvm::formatv(structuredOpIteratorTypesFormat, className,
 823                         iteratorsStr);
 824   } else {
 825     os << llvm::formatv(rankPolyStructuredOpIteratorTypesFormat, className);
 826   }
 827
 828   // Generating the getIndexingMaps() method.
 829   if (auto &staticMaps =
 830           opConfig.structuredOp->indexingMaps.staticIndexingMaps) {
 831     if (staticMaps->empty())
 832       return emitError(genContext.getLoc()) << "op has no indexing maps";
 833     if (!isRankPolymorphic) {
 834       AffineMap firstMap = staticMaps->front().affineMap();
 835
 836       // Symbol bindings.
 837       {
 838         // For each symbol, generate a declaration for it, either with an
 839         // AffineSymbolExpr or an AffineConstantExpr (if the symbol derives from
 840         // an attribute).
 841         // TODO: Possibly lift into a top-level method.
 842         static const char structuredOpSymbolBindingsFormat[] = R"FMT(
 843 static SmallVector<AffineExpr> getSymbolBindings({0} self) {
 844   MLIRContext *context = self.getContext();
 845   SmallVector<AffineExpr> exprs;
 846 {1}
 847   return exprs;
 848 }
 849 )FMT";
 850
 851         unsigned symbolCount = firstMap.getNumSymbols();
 852         SmallVector<std::string> symbolBindings;
 853         for (unsigned i = 0; i < symbolCount; ++i) {
 854           symbolBindings.push_back(llvm::formatv(
 855               "  exprs.push_back(getAffineSymbolExpr({0}, context));", i));
 856         }
 857
 858         // Access an index attribute. Parameters:
 859         // {0}: Attribute name
 860         // {1}: Symbol position
 861         // {2}: Attribute index
 862         static const char structuredOpAccessAttrFormat[] = R"FMT(
 863 int64_t cst{1} = self.get{0}().getValues<int64_t>()[{2}];
 864 exprs.push_back(getAffineConstantExpr(cst{1}, context));
 865 )FMT";
 866         // Update all symbol bindings mapped to an attribute.
 867         for (LinalgOperandDef &arg : opConfig.structuredOp->args) {
 868           if (arg.kind != LinalgOperandDefKind::IndexAttr)
 869             continue;
 870           assert(arg.indexAttrMap);
 871           for (auto [idx, result] :
 872                llvm::enumerate(arg.indexAttrMap->affineMap().getResults())) {
 873             if (auto symbol = result.dyn_cast<AffineSymbolExpr>()) {
 874               std::string argName = arg.name;
 875               argName[0] = toupper(argName[0]);
 876               symbolBindings[symbol.getPosition()] =
 877                   llvm::formatv(structuredOpAccessAttrFormat, argName,
 878                                 symbol.getPosition(), idx);
 879             }
 880           }
 881         }
 882
 883         std::string symbolBindingsStr;
 884         llvm::raw_string_ostream symbolBindingsSs(symbolBindingsStr);
 885         llvm::interleave(symbolBindings, symbolBindingsSs, "\n");
 886         symbolBindingsSs.flush();
 887
 888         os << llvm::formatv(structuredOpSymbolBindingsFormat, className,
 889                             symbolBindingsStr);
 890       }
 891
 892       // Indexing maps.
 893       {
 894         unsigned dimCount = firstMap.getNumDims();
 895
 896         // Generate a comma-separated list of dim identifiers to be passed to
 897         // bindDims, ensuring tht AffineExpr identifiers are bound in the right
 898         // order to the proper AffineDimExpr.
 899         // This results in vars in scope like: d0, d1, d2...
 900         SmallVector<unsigned> dimIndices;
 901         for (unsigned i = 0; i < dimCount; ++i)
 902           dimIndices.push_back(i);
 903         std::string dimIdentsStr;
 904         llvm::raw_string_ostream dimIdentsSs(dimIdentsStr);
 905         llvm::interleaveComma(dimIndices, dimIdentsSs,
 906                               [&](unsigned i) { dimIdentsSs << "d" << i; });
 907         dimIdentsSs.flush();
 908
 909         // Statements to add and simplify each affine map.
 910         SmallVector<std::string> stmts;
 911         for (auto &indexingMap : *staticMaps) {
 912           // TODO: Assert that dim and symbol count match the first.
 913           stmts.push_back(
 914               llvm::formatv("maps.push_back({0});",
 915                             generateCppExpression(indexingMap, "context")));
 916           stmts.push_back(llvm::formatv(
 917               "maps.back() = "
 918               "simplifyAffineMap(maps.back().replaceDimsAndSymbols({{}, "
 919               "symbolBindings, {0}, 0));",
 920               dimCount));
 921         }
 922
 923         // TODO: This needs to be memoized and/or converted to non-parser based
 924         // C++ codegen prior to real use.
 925         os << llvm::formatv(structuredOpIndexingMapsFormat, className,
 926                             dimIdentsStr, interleaveToString(stmts, "\n  "));
 927       }
 928     } else {
 929       os << llvm::formatv(rankPolyStructuredOpIndexingMapsFormat, className);
 930     }
 931   } else {
 932     return emitError(genContext.getLoc())
 933            << "generating code for non static indexing maps not currently "
 934               "supported";
 935   }
 936
 937   // getNumRegionArgs()
 938   {
 939     // Generates a getNumRegionArgs() method. Parameters:
 940     // {0}: Class name
 941     // {1}: Number of region args
 942     static const char structuredOpGetNumRegionArgsFormat[] = R"FMT(
 943 unsigned {0}::getNumRegionArgs() {{ return {1}; }
 944 )FMT";
 945     os << llvm::formatv(structuredOpGetNumRegionArgsFormat, className,
 946                         numOfArgs);
 947   }
 948
 949   // getLibraryCallName()
 950   {
 951     // Generates a getLibraryCallName method. Parameters:
 952     // {0}: Class name
 953     static const char structuredOpGetLibraryCallFormat[] = R"FMT(
 954 std::string {0}::getLibraryCallName() {{
 955   return generateLibraryCallName(getOperation());
 956 }
 957 )FMT";
 958     os << llvm::formatv(structuredOpGetLibraryCallFormat, className);
 959   }
 960
 961   // hasDynamicIndexingMaps() and verifyIndexingMapRequiredAttributes()
 962   if (llvm::any_of(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
 963         return arg.kind == LinalgOperandDefKind::IndexAttr;
 964       })) {
 965     std::vector<std::string> attrVerifications;
 966     for (LinalgOperandDef &arg : opConfig.structuredOp->args) {
 967       if (arg.kind != LinalgOperandDefKind::IndexAttr)
 968         continue;
 969       assert(arg.indexAttrMap);
 970       // Verify index attribute. Paramters:
 971       // {0}: Attribute name
 972       // {1}: Attribute size
 973       static const char attrFmt[] = R"FMT(
 974 if (auto attr = op->getAttrOfType<DenseElementsAttr>("{0}")) {{
 975   if (!attr.getType().getElementType().isInteger(64))
 976     return op->emitError("incorrect element type for index attribute '{0}'");
 977   if (attr.getType().getShape() != ArrayRef<int64_t>{{ {1} })
 978     return op->emitError("incorrect shape for index attribute '{0}'");
 979 }
 980 )FMT";
 981       attrVerifications.push_back(llvm::formatv(
 982           attrFmt, arg.name, arg.indexAttrMap->affineMap().getNumResults()));
 983     }
 984
 985     // Generates the verifyIndexingMapRequiredAttributes method. Parameters:
 986     // {0}: Class name
 987     // {1}: Attribute verification
 988     static const char structuredOpVerifyIndexingMapRequiredAttributes[] = R"FMT(
 989 bool {0}::hasDynamicIndexingMaps() {{ return true; }
 990 LogicalResult {0}::verifyIndexingMapRequiredAttributes() {{
 991   Operation *op = getOperation();
 992   {1}
 993   return success();
 994 }
 995 )FMT";
 996     os << llvm::formatv(structuredOpVerifyIndexingMapRequiredAttributes,
 997                         className, llvm::join(attrVerifications, "\n"));
 998   }
 999
1000   // regionBuilder()
1001   {
1002     // Generates a regionBuilder method. Parameters.
1003     // {0}: Class name
1004     // {1}: Number of args
1005     // {2}: Attributes
1006     // {3}: Statements
1007     static const char structuredOpRegionBuilderFormat[] = R"FMT(
1008 void {0}::regionBuilder(ImplicitLocOpBuilder &b,
1009                         Block &block, ArrayRef<NamedAttribute> attrs) {{
1010   assert({1} > 0 && block.getNumArguments() == {1} &&
1011          "{0} regionBuilder expects {1} (>=0) args");
1012   RegionBuilderHelper helper(block.getArgument(0).getContext(), block);
1013   SmallVector<Value> yields;
1014   {2}
1015   {3}
1016   helper.yieldOutputs(yields);
1017 }
1018 )FMT";
1019     auto &args = opConfig.structuredOp->args;
1020     auto &assignments = opConfig.structuredOp->assignments;
1021     size_t generatedAssignmentCount = 0;
1022     int localCounter = 0;
1023     SmallVector<std::string> attrs;
1024     SmallVector<std::string> stmts;
1025     for (LinalgOperandDef &arg : args) {
1026       if (!isFunctionAttribute(arg.kind))
1027         continue;
1028       // Obtain the type function attribute values. Parameters.
1029       // {0}: enum name
1030       // {1}: attribute name
1031       // {2}: default type function name
1032       static const char attrDef[] = R"FMT(
1033 {0} {1}Val = {0}::{2};
1034 auto {1}Iter = llvm::find_if(attrs, [&](const NamedAttribute &attr) {{
1035                               return attr.getName() == "{1}"; });
1036 if ({1}Iter != attrs.end()) {{
1037   if (auto attr = llvm::dyn_cast<{0}Attr>({1}Iter->getValue()))
1038     {1}Val = attr.getValue();
1039 }
1040 )FMT";
1041       std::string enumName = convertOperandKindToEnumName(arg.kind);
1042       attrs.push_back(
1043           llvm::formatv(attrDef, enumName, arg.name, arg.defaultFn));
1044     }
1045     for (LinalgOperandDef &arg : args) {
1046       if (arg.kind != LinalgOperandDefKind::OutputTensor)
1047         continue;
1048
1049       // Find the assignment that correlates with the argument.
1050       ScalarAssign *assignment = findAssignment(arg.name, assignments);
1051       if (!assignment)
1052         return emitError(genContext.getLoc())
1053                << "no assignment found for output argument " << arg.name;
1054       ++generatedAssignmentCount;
1055
1056       // Recursively generate the expression.
1057       std::function<std::optional<std::string>(ScalarExpression &)>
1058           generateExpression =
1059               [&](ScalarExpression &expression) -> std::optional<std::string> {
1060         if (expression.arg) {
1061           // Argument reference.
1062           std::optional<int> argIndex =
1063               findTensorDefArgIndex(*expression.arg, args);
1064           if (!argIndex) {
1065             emitError(genContext.getLoc())
1066                 << "scalar argument not defined on the op: " << *expression.arg;
1067             return std::nullopt;
1068           }
1069           return std::string(
1070               llvm::formatv("block.getArgument({0})", *argIndex));
1071         }
1072         if (expression.constant) {
1073           std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
1074           stmts.push_back(
1075               llvm::formatv(R"FMT(Value {0} = helper.constant("{1}");)FMT",
1076                             cppIdent, expression.constant));
1077           return cppIdent;
1078         }
1079         if (expression.index) {
1080           // Access an iteration index.
1081           std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
1082           stmts.push_back(llvm::formatv("Value {0} = helper.index({1});",
1083                                         cppIdent, *expression.index));
1084           return cppIdent;
1085         }
1086         if (expression.scalarFn) {
1087           std::string enumName =
1088               convertFunctionKindToEnumName(expression.scalarFn->kind);
1089
1090           // Get the function or attribute name.
1091           assert(expression.scalarFn->fnName || expression.scalarFn->attrName);
1092           std::string funcType;
1093           if (expression.scalarFn->fnName) {
1094             funcType = llvm::formatv("{0}::{1}", enumName,
1095                                      *expression.scalarFn->fnName);
1096           }
1097           if (expression.scalarFn->attrName) {
1098             if (llvm::none_of(args, [&](LinalgOperandDef &arg) {
1099                   return isFunctionAttribute(arg.kind) &&
1100                          arg.name == *expression.scalarFn->attrName;
1101                 })) {
1102               emitError(genContext.getLoc()) << "missing function attribute "
1103                                              << *expression.scalarFn->attrName;
1104             }
1105             funcType = llvm::formatv("{0}Val", *expression.scalarFn->attrName);
1106           }
1107           assert(!funcType.empty());
1108
1109           // Add the optional type parameter to the operands.
1110           SmallVector<std::string> operandCppValues;
1111           if (expression.scalarFn->kind == ScalarFnKind::Type) {
1112             assert(expression.scalarFn->typeVar.has_value());
1113             std::optional<std::string> typeCppValue =
1114                 findTypeValue(*expression.scalarFn->typeVar, args);
1115             if (!typeCppValue) {
1116               emitError(genContext.getLoc())
1117                   << "type variable " << *expression.scalarFn->typeVar
1118                   << ", used in a type conversion, must map to a predefined or "
1119                   << "an argument type but it does not";
1120               return std::nullopt;
1121             }
1122             operandCppValues.push_back(*typeCppValue);
1123           }
1124
1125           // Collect the scalar operands.
1126           for (ScalarExpression &operand : expression.scalarFn->operands) {
1127             auto operandCppValue = generateExpression(operand);
1128             if (!operandCppValue)
1129               return std::nullopt;
1130             operandCppValues.push_back(*operandCppValue);
1131           }
1132
1133           // Call the function builder.
1134           std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
1135           stmts.push_back(llvm::formatv(
1136               "Value {0} = helper.build{1}({2}, {3});", cppIdent, enumName,
1137               funcType, interleaveToString(operandCppValues, ", ")));
1138           return cppIdent;
1139         }
1140         emitError(genContext.getLoc()) << "unknown ScalarExpression type";
1141         return std::nullopt;
1142       };
1143       std::optional<std::string> cppValue =
1144           generateExpression(assignment->value);
1145       if (!cppValue)
1146         return failure();
1147       stmts.push_back(llvm::formatv("yields.push_back({0});", *cppValue));
1148     }
1149
1150     if (generatedAssignmentCount != assignments.size())
1151       return emitError(genContext.getLoc())
1152              << "mismatched number of assignments vs output arguments";
1153
1154     os << llvm::formatv(structuredOpRegionBuilderFormat, className, numOfArgs,
1155                         interleaveToString(attrs, "\n  "),
1156                         interleaveToString(stmts, "\n  "));
1157   }
1158
1159   // Parser and printer.
1160   os << llvm::formatv(structuredOpParserFormat, className);
1161
1162   // Canonicalizers and folders.
1163   os << llvm::formatv(structuredOpFoldersFormat, className);
1164
1165   return success();
1166 }
1167
1168 static LogicalResult generateOp(LinalgOpConfig &opConfig,
1169                                 GenerationContext &genContext) {
1170   // Switch on op type being generated.
1171   if (opConfig.structuredOp) {
1172     return success(
1173         succeeded(generateNamedGenericOpOds(opConfig, genContext)) &&
1174         succeeded(generateNamedGenericOpDefns(opConfig, genContext)));
1175   }
1176   return emitError(genContext.getLoc()) << "unsupported operation type";
1177 }
1178
1179 //===----------------------------------------------------------------------===//
1180 // Command line options and main
1181 //===----------------------------------------------------------------------===//
1182
1183 static llvm::cl::opt<std::string>
1184     inputFilename(llvm::cl::Positional, llvm::cl::desc("<input file>"),
1185                   llvm::cl::init("-"), llvm::cl::value_desc("YAML filename"));
1186
1187 static llvm::cl::opt<std::string>
1188     outputOdsDeclFilename("o-ods-decl", llvm::cl::desc("ODS output filename"),
1189                           llvm::cl::value_desc("filename"), llvm::cl::init(""));
1190
1191 static llvm::cl::opt<std::string>
1192     outputCppImplFilename("o-impl",
1193                           llvm::cl::desc("C++ implementation file name"),
1194                           llvm::cl::value_desc("filename"), llvm::cl::init(""));
1195
1196 int main(int argc, char **argv) {
1197   llvm::cl::ParseCommandLineOptions(argc, argv, "Linalg ODS Gen from YAML");
1198
1199   // Set up the input file.
1200   std::string errorMessage;
1201   std::unique_ptr<llvm::MemoryBuffer> file =
1202       mlir::openInputFile(inputFilename, &errorMessage);
1203   if (!file) {
1204     llvm::errs() << errorMessage << "\n";
1205     return 1;
1206   }
1207
1208   MLIRContext mlirContext;
1209   LinalgYAMLContext yamlContext{&mlirContext};
1210
1211   std::vector<LinalgOpConfig> opConfigs;
1212
1213   // Parse input.
1214   Input yin(file->getBuffer(), &yamlContext);
1215   yin >> opConfigs;
1216
1217   if (yin.error())
1218     return 1;
1219
1220   // Open output files.
1221   std::unique_ptr<llvm::ToolOutputFile> outputOdsDecl;
1222   if (!outputOdsDeclFilename.empty()) {
1223     outputOdsDecl = openOutputFile(outputOdsDeclFilename, &errorMessage);
1224     if (!outputOdsDecl) {
1225       llvm::errs() << errorMessage << "\n";
1226       return 1;
1227     }
1228   }
1229
1230   std::unique_ptr<llvm::ToolOutputFile> outputCppImpl;
1231   if (!outputCppImplFilename.empty()) {
1232     outputCppImpl = openOutputFile(outputCppImplFilename, &errorMessage);
1233     if (!outputCppImpl) {
1234       llvm::errs() << errorMessage << "\n";
1235       return 1;
1236     }
1237   }
1238
1239   if (!outputOdsDecl && !outputCppImpl) {
1240     llvm::errs() << "error: No output files specified\n";
1241     return 1;
1242   }
1243
1244   // Generate.
1245   GenerationContext genContext(&mlirContext,
1246                                outputOdsDecl ? &outputOdsDecl->os() : nullptr,
1247                                outputCppImpl ? &outputCppImpl->os() : nullptr);
1248
1249   for (auto &opConfig : opConfigs) {
1250     if (!opConfig.metadata) {
1251       emitError(genContext.getLoc())
1252           << "missing operation metadata on subsequent op";
1253       return 1;
1254     }
1255
1256     genContext.setLoc(NameLoc::get(
1257         StringAttr::get(&mlirContext, opConfig.metadata->cppClassName)));
1258     if (failed(generateOp(opConfig, genContext))) {
1259       return 1;
1260     }
1261   }
1262
1263   if (outputOdsDecl)
1264     outputOdsDecl->keep();
1265   if (outputCppImpl)
1266     outputCppImpl->keep();
1267
1268   return 0;
1269 }