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(mlir::parseAttribute("{0}", {1}).cast<AffineMapAttr>().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 - 1, 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 = [{
 698     {1}
 699   }];
 700 )FMT";
 701     StringRef summary, description;
 702     std::tie(summary, description) =
 703         StringRef(*opConfig.metadata->doc).trim().split('\n');
 704     doc = llvm::formatv(structuredOpDocFmt, summary.trim(), description.trim());
 705   }
 706
 707   interfaceNameList = interleaveToString(opConfig.metadata->implements, ", ");
 708
 709   std::string definitionList;
 710   for (const std::string &definition : opConfig.metadata->defines) {
 711     static const char definitionFmt[] = "let {0} = 1;\n";
 712     definitionList.append(llvm::formatv(definitionFmt, definition));
 713   }
 714
 715   if (llvm::any_of(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
 716         return isAttribute(arg.kind);
 717       })) {
 718     SmallVector<std::string> attrDefs;
 719     SmallVector<std::string> attrParams;
 720     SmallVector<std::string> attrStmts;
 721     for (LinalgOperandDef &arg : opConfig.structuredOp->args) {
 722       static const char paramFmt[] = "\"Attribute\":${0}";
 723       static const char stmtFmt[] = "$_state.addAttribute(\"{0}\", {0});";
 724       // Add the type conversion attributes to the op definition and builders.
 725       if (isFunctionAttribute(arg.kind)) {
 726         assert(arg.defaultFn);
 727         std::string enumName = convertOperandKindToEnumName(arg.kind);
 728         static const char typeFmt[] = "{0}::{1}";
 729         static const char defFmt[] =
 730             "DefaultValuedOptionalAttr<{0}, \"{1}\">:${2}";
 731         attrDefs.push_back(llvm::formatv(
 732             defFmt, llvm::formatv("{0}Attr", enumName),
 733             llvm::formatv(typeFmt, enumName, arg.defaultFn), arg.name));
 734         attrParams.push_back(llvm::formatv(paramFmt, arg.name));
 735         attrStmts.push_back(llvm::formatv(stmtFmt, arg.name));
 736       }
 737       // Add the index attributes to the op definition and builders.
 738       if (arg.kind == LinalgOperandDefKind::IndexAttr) {
 739         assert(arg.indexAttrMap.has_value());
 740         assert(arg.defaultIndices.has_value());
 741         size_t size = arg.indexAttrMap->affineMap().getNumResults();
 742         assert(arg.defaultIndices->size() == size);
 743         static const char typeFmt[] = "RankedI64ElementsAttr<[{0}]>";
 744         static const char defFmt[] =
 745             "DefaultValuedOptionalAttr<{0}, \"{ {1} }\">:${2}";
 746         std::string defaultVals;
 747         llvm::raw_string_ostream ss(defaultVals);
 748         llvm::interleave(
 749             *arg.defaultIndices, ss,
 750             [&](int64_t val) { ss << "static_cast<int64_t>(" << val << ")"; },
 751             ", ");
 752         attrDefs.push_back(llvm::formatv(defFmt, llvm::formatv(typeFmt, size),
 753                                          ss.str(), arg.name));
 754         attrParams.push_back(llvm::formatv(paramFmt, arg.name));
 755         attrStmts.push_back(llvm::formatv(stmtFmt, arg.name));
 756       }
 757     }
 758     if (llvm::any_of(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
 759           return arg.kind == LinalgOperandDefKind::IndexAttr;
 760         })) {
 761       attrMethods = R"(
 762         bool hasDynamicIndexingMaps();
 763         LogicalResult verifyIndexingMapRequiredAttributes();
 764       )";
 765     }
 766     attrList = ",\n" + llvm::join(attrDefs, ",\n");
 767     attrBuilder = llvm::formatv(
 768         structuredOpBuilderFormat, opConfig.metadata->cppClassName,
 769         llvm::join(attrParams, ", "), llvm::join(attrStmts, "\n"));
 770   }
 771
 772   os << llvm::formatv(structuredOpOdsHeaderFormat,
 773                       opConfig.metadata->cppClassName, opConfig.metadata->name,
 774                       interfaceNameList, doc, attrList, attrBuilder,
 775                       definitionList, attrMethods);
 776
 777   return success();
 778 }
 779
 780 static LogicalResult
 781 generateNamedGenericOpDefns(LinalgOpConfig &opConfig,
 782                             GenerationContext &genContext) {
 783   if (!genContext.shouldGenerateDefns())
 784     return success();
 785
 786   raw_ostream &os = genContext.defns();
 787   StringRef className = opConfig.metadata->cppClassName;
 788
 789   // Implementation banner.
 790   std::string bannerComment = llvm::formatv("Implementation of {0}", className);
 791   os << llvm::formatv(bannerFormat, bannerComment);
 792
 793   // Compute the number of scalar and tensor arguments.
 794   int64_t numOfArgs =
 795       llvm::count_if(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
 796         return arg.kind == LinalgOperandDefKind::InputTensor ||
 797                arg.kind == LinalgOperandDefKind::Scalar ||
 798                arg.kind == LinalgOperandDefKind::OutputTensor;
 799       });
 800
 801   // An operation that accesses only scalars and scalar/rank zero tensors is
 802   // rank polymorhpic. We implement rank polymorphism by generating different
 803   // indexing maps and iterators that match the rank of the first output tensor.
 804   // An operation is rank polymorphic if the iteration domain has rank zero.
 805   bool isRankPolymorphic = opConfig.structuredOp->iteratorTypes.empty();
 806
 807   // Generate the iterator_types() method.
 808   if (!isRankPolymorphic) {
 809     std::string iteratorsStr;
 810     llvm::raw_string_ostream ss(iteratorsStr);
 811     llvm::interleaveComma(opConfig.structuredOp->iteratorTypes, ss,
 812                           [&](LinalgIteratorTypeDef it) {
 813                             switch (it) {
 814                             case LinalgIteratorTypeDef::parallel:
 815                               ss << "utils::IteratorType::parallel";
 816                               break;
 817                             case LinalgIteratorTypeDef::reduction:
 818                               ss << "utils::IteratorType::reduction";
 819                               break;
 820                             }
 821                           });
 822     ss.flush();
 823     os << llvm::formatv(structuredOpIteratorTypesFormat, className,
 824                         iteratorsStr);
 825   } else {
 826     os << llvm::formatv(rankPolyStructuredOpIteratorTypesFormat, className);
 827   }
 828
 829   // Generating the getIndexingMaps() method.
 830   if (auto &staticMaps =
 831           opConfig.structuredOp->indexingMaps.staticIndexingMaps) {
 832     if (staticMaps->empty())
 833       return emitError(genContext.getLoc()) << "op has no indexing maps";
 834     if (!isRankPolymorphic) {
 835       AffineMap firstMap = staticMaps->front().affineMap();
 836
 837       // Symbol bindings.
 838       {
 839         // For each symbol, generate a declaration for it, either with an
 840         // AffineSymbolExpr or an AffineConstantExpr (if the symbol derives from
 841         // an attribute).
 842         // TODO: Possibly lift into a top-level method.
 843         static const char structuredOpSymbolBindingsFormat[] = R"FMT(
 844 static SmallVector<AffineExpr> getSymbolBindings({0} self) {
 845   MLIRContext *context = self.getContext();
 846   SmallVector<AffineExpr> exprs;
 847 {1}
 848   return exprs;
 849 }
 850 )FMT";
 851
 852         unsigned symbolCount = firstMap.getNumSymbols();
 853         SmallVector<std::string> symbolBindings;
 854         for (unsigned i = 0; i < symbolCount; ++i) {
 855           symbolBindings.push_back(llvm::formatv(
 856               "  exprs.push_back(getAffineSymbolExpr({0}, context));", i));
 857         }
 858
 859         // Access an index attribute. Parameters:
 860         // {0}: Attribute name
 861         // {1}: Symbol position
 862         // {2}: Attribute index
 863         static const char structuredOpAccessAttrFormat[] = R"FMT(
 864 int64_t cst{1} = self.get{0}().getValues<int64_t>()[{2}];
 865 exprs.push_back(getAffineConstantExpr(cst{1}, context));
 866 )FMT";
 867         // Update all symbol bindings mapped to an attribute.
 868         for (LinalgOperandDef &arg : opConfig.structuredOp->args) {
 869           if (arg.kind != LinalgOperandDefKind::IndexAttr)
 870             continue;
 871           assert(arg.indexAttrMap);
 872           for (auto [idx, result] :
 873                llvm::enumerate(arg.indexAttrMap->affineMap().getResults())) {
 874             if (auto symbol = result.dyn_cast<AffineSymbolExpr>()) {
 875               std::string argName = arg.name;
 876               argName[0] = toupper(argName[0]);
 877               symbolBindings[symbol.getPosition()] =
 878                   llvm::formatv(structuredOpAccessAttrFormat, argName,
 879                                 symbol.getPosition(), idx);
 880             }
 881           }
 882         }
 883
 884         std::string symbolBindingsStr;
 885         llvm::raw_string_ostream symbolBindingsSs(symbolBindingsStr);
 886         llvm::interleave(symbolBindings, symbolBindingsSs, "\n");
 887         symbolBindingsSs.flush();
 888
 889         os << llvm::formatv(structuredOpSymbolBindingsFormat, className,
 890                             symbolBindingsStr);
 891       }
 892
 893       // Indexing maps.
 894       {
 895         unsigned dimCount = firstMap.getNumDims();
 896
 897         // Generate a comma-separated list of dim identifiers to be passed to
 898         // bindDims, ensuring tht AffineExpr identifiers are bound in the right
 899         // order to the proper AffineDimExpr.
 900         // This results in vars in scope like: d0, d1, d2...
 901         SmallVector<unsigned> dimIndices;
 902         for (unsigned i = 0; i < dimCount; ++i)
 903           dimIndices.push_back(i);
 904         std::string dimIdentsStr;
 905         llvm::raw_string_ostream dimIdentsSs(dimIdentsStr);
 906         llvm::interleaveComma(dimIndices, dimIdentsSs,
 907                               [&](unsigned i) { dimIdentsSs << "d" << i; });
 908         dimIdentsSs.flush();
 909
 910         // Statements to add and simplify each affine map.
 911         SmallVector<std::string> stmts;
 912         for (auto &indexingMap : *staticMaps) {
 913           // TODO: Assert that dim and symbol count match the first.
 914           stmts.push_back(
 915               llvm::formatv("maps.push_back({0});",
 916                             generateCppExpression(indexingMap, "context")));
 917           stmts.push_back(llvm::formatv(
 918               "maps.back() = "
 919               "simplifyAffineMap(maps.back().replaceDimsAndSymbols({{}, "
 920               "symbolBindings, {0}, 0));",
 921               dimCount));
 922         }
 923
 924         // TODO: This needs to be memoized and/or converted to non-parser based
 925         // C++ codegen prior to real use.
 926         os << llvm::formatv(structuredOpIndexingMapsFormat, className,
 927                             dimIdentsStr, interleaveToString(stmts, "\n  "));
 928       }
 929     } else {
 930       os << llvm::formatv(rankPolyStructuredOpIndexingMapsFormat, className);
 931     }
 932   } else {
 933     return emitError(genContext.getLoc())
 934            << "generating code for non static indexing maps not currently "
 935               "supported";
 936   }
 937
 938   // getNumRegionArgs()
 939   {
 940     // Generates a getNumRegionArgs() method. Parameters:
 941     // {0}: Class name
 942     // {1}: Number of region args
 943     static const char structuredOpGetNumRegionArgsFormat[] = R"FMT(
 944 unsigned {0}::getNumRegionArgs() {{ return {1}; }
 945 )FMT";
 946     os << llvm::formatv(structuredOpGetNumRegionArgsFormat, className,
 947                         numOfArgs);
 948   }
 949
 950   // getLibraryCallName()
 951   {
 952     // Generates a getLibraryCallName method. Parameters:
 953     // {0}: Class name
 954     static const char structuredOpGetLibraryCallFormat[] = R"FMT(
 955 std::string {0}::getLibraryCallName() {{
 956   return generateLibraryCallName(getOperation());
 957 }
 958 )FMT";
 959     os << llvm::formatv(structuredOpGetLibraryCallFormat, className);
 960   }
 961
 962   // hasDynamicIndexingMaps() and verifyIndexingMapRequiredAttributes()
 963   if (llvm::any_of(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
 964         return arg.kind == LinalgOperandDefKind::IndexAttr;
 965       })) {
 966     std::vector<std::string> attrVerifications;
 967     for (LinalgOperandDef &arg : opConfig.structuredOp->args) {
 968       if (arg.kind != LinalgOperandDefKind::IndexAttr)
 969         continue;
 970       assert(arg.indexAttrMap);
 971       // Verify index attribute. Paramters:
 972       // {0}: Attribute name
 973       // {1}: Attribute size
 974       static const char attrFmt[] = R"FMT(
 975 if (auto attr = op->getAttrOfType<DenseElementsAttr>("{0}")) {{
 976   if (!attr.getType().getElementType().isInteger(64))
 977     return op->emitError("incorrect element type for index attribute '{0}'");
 978   if (attr.getType().getShape() != ArrayRef<int64_t>{{ {1} })
 979     return op->emitError("incorrect shape for index attribute '{0}'");
 980 }
 981 )FMT";
 982       attrVerifications.push_back(llvm::formatv(
 983           attrFmt, arg.name, arg.indexAttrMap->affineMap().getNumResults()));
 984     }
 985
 986     // Generates the verifyIndexingMapRequiredAttributes method. Parameters:
 987     // {0}: Class name
 988     // {1}: Attribute verification
 989     static const char structuredOpVerifyIndexingMapRequiredAttributes[] = R"FMT(
 990 bool {0}::hasDynamicIndexingMaps() {{ return true; }
 991 LogicalResult {0}::verifyIndexingMapRequiredAttributes() {{
 992   Operation *op = getOperation();
 993   {1}
 994   return success();
 995 }
 996 )FMT";
 997     os << llvm::formatv(structuredOpVerifyIndexingMapRequiredAttributes,
 998                         className, llvm::join(attrVerifications, "\n"));
 999   }
1000
1001   // regionBuilder()
1002   {
1003     // Generates a regionBuilder method. Parameters.
1004     // {0}: Class name
1005     // {1}: Number of args
1006     // {2}: Attributes
1007     // {3}: Statements
1008     static const char structuredOpRegionBuilderFormat[] = R"FMT(
1009 void {0}::regionBuilder(ImplicitLocOpBuilder &b,
1010                         Block &block, ArrayRef<NamedAttribute> attrs) {{
1011   assert({1} > 0 && block.getNumArguments() == {1} &&
1012          "{0} regionBuilder expects {1} (>=0) args");
1013   RegionBuilderHelper helper(block.getArgument(0).getContext(), block);
1014   SmallVector<Value> yields;
1015   {2}
1016   {3}
1017   helper.yieldOutputs(yields);
1018 }
1019 )FMT";
1020     auto &args = opConfig.structuredOp->args;
1021     auto &assignments = opConfig.structuredOp->assignments;
1022     size_t generatedAssignmentCount = 0;
1023     int localCounter = 0;
1024     SmallVector<std::string> attrs;
1025     SmallVector<std::string> stmts;
1026     for (LinalgOperandDef &arg : args) {
1027       if (!isFunctionAttribute(arg.kind))
1028         continue;
1029       // Obtain the type function attribute values. Parameters.
1030       // {0}: enum name
1031       // {1}: attribute name
1032       // {2}: default type function name
1033       static const char attrDef[] = R"FMT(
1034 {0} {1}Val = {0}::{2};
1035 auto {1}Iter = llvm::find_if(attrs, [&](const NamedAttribute &attr) {{
1036                               return attr.getName() == "{1}"; });
1037 if ({1}Iter != attrs.end()) {{
1038   if (auto attr = {1}Iter->getValue().dyn_cast<{0}Attr>())
1039     {1}Val = attr.getValue();
1040 }
1041 )FMT";
1042       std::string enumName = convertOperandKindToEnumName(arg.kind);
1043       attrs.push_back(
1044           llvm::formatv(attrDef, enumName, arg.name, arg.defaultFn));
1045     }
1046     for (LinalgOperandDef &arg : args) {
1047       if (arg.kind != LinalgOperandDefKind::OutputTensor)
1048         continue;
1049
1050       // Find the assignment that correlates with the argument.
1051       ScalarAssign *assignment = findAssignment(arg.name, assignments);
1052       if (!assignment)
1053         return emitError(genContext.getLoc())
1054                << "no assignment found for output argument " << arg.name;
1055       ++generatedAssignmentCount;
1056
1057       // Recursively generate the expression.
1058       std::function<std::optional<std::string>(ScalarExpression &)>
1059           generateExpression =
1060               [&](ScalarExpression &expression) -> std::optional<std::string> {
1061         if (expression.arg) {
1062           // Argument reference.
1063           std::optional<int> argIndex =
1064               findTensorDefArgIndex(*expression.arg, args);
1065           if (!argIndex) {
1066             emitError(genContext.getLoc())
1067                 << "scalar argument not defined on the op: " << *expression.arg;
1068             return std::nullopt;
1069           }
1070           return std::string(
1071               llvm::formatv("block.getArgument({0})", *argIndex));
1072         }
1073         if (expression.constant) {
1074           std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
1075           stmts.push_back(
1076               llvm::formatv(R"FMT(Value {0} = helper.constant("{1}");)FMT",
1077                             cppIdent, expression.constant));
1078           return cppIdent;
1079         }
1080         if (expression.index) {
1081           // Access an iteration index.
1082           std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
1083           stmts.push_back(llvm::formatv("Value {0} = helper.index({1});",
1084                                         cppIdent, *expression.index));
1085           return cppIdent;
1086         }
1087         if (expression.scalarFn) {
1088           std::string enumName =
1089               convertFunctionKindToEnumName(expression.scalarFn->kind);
1090
1091           // Get the function or attribute name.
1092           assert(expression.scalarFn->fnName || expression.scalarFn->attrName);
1093           std::string funcType;
1094           if (expression.scalarFn->fnName) {
1095             funcType = llvm::formatv("{0}::{1}", enumName,
1096                                      *expression.scalarFn->fnName);
1097           }
1098           if (expression.scalarFn->attrName) {
1099             if (llvm::none_of(args, [&](LinalgOperandDef &arg) {
1100                   return isFunctionAttribute(arg.kind) &&
1101                          arg.name == *expression.scalarFn->attrName;
1102                 })) {
1103               emitError(genContext.getLoc()) << "missing function attribute "
1104                                              << *expression.scalarFn->attrName;
1105             }
1106             funcType = llvm::formatv("{0}Val", *expression.scalarFn->attrName);
1107           }
1108           assert(!funcType.empty());
1109
1110           // Add the optional type parameter to the operands.
1111           SmallVector<std::string> operandCppValues;
1112           if (expression.scalarFn->kind == ScalarFnKind::Type) {
1113             assert(expression.scalarFn->typeVar.has_value());
1114             std::optional<std::string> typeCppValue =
1115                 findTypeValue(*expression.scalarFn->typeVar, args);
1116             if (!typeCppValue) {
1117               emitError(genContext.getLoc())
1118                   << "type variable " << *expression.scalarFn->typeVar
1119                   << ", used in a type conversion, must map to a predefined or "
1120                   << "an argument type but it does not";
1121               return std::nullopt;
1122             }
1123             operandCppValues.push_back(*typeCppValue);
1124           }
1125
1126           // Collect the scalar operands.
1127           for (ScalarExpression &operand : expression.scalarFn->operands) {
1128             auto operandCppValue = generateExpression(operand);
1129             if (!operandCppValue)
1130               return std::nullopt;
1131             operandCppValues.push_back(*operandCppValue);
1132           }
1133
1134           // Call the function builder.
1135           std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
1136           stmts.push_back(llvm::formatv(
1137               "Value {0} = helper.build{1}({2}, {3});", cppIdent, enumName,
1138               funcType, interleaveToString(operandCppValues, ", ")));
1139           return cppIdent;
1140         }
1141         emitError(genContext.getLoc()) << "unknown ScalarExpression type";
1142         return std::nullopt;
1143       };
1144       std::optional<std::string> cppValue =
1145           generateExpression(assignment->value);
1146       if (!cppValue)
1147         return failure();
1148       stmts.push_back(llvm::formatv("yields.push_back({0});", *cppValue));
1149     }
1150
1151     if (generatedAssignmentCount != assignments.size())
1152       return emitError(genContext.getLoc())
1153              << "mismatched number of assignments vs output arguments";
1154
1155     os << llvm::formatv(structuredOpRegionBuilderFormat, className, numOfArgs,
1156                         interleaveToString(attrs, "\n  "),
1157                         interleaveToString(stmts, "\n  "));
1158   }
1159
1160   // Parser and printer.
1161   os << llvm::formatv(structuredOpParserFormat, className);
1162
1163   // Canonicalizers and folders.
1164   os << llvm::formatv(structuredOpFoldersFormat, className);
1165
1166   return success();
1167 }
1168
1169 static LogicalResult generateOp(LinalgOpConfig &opConfig,
1170                                 GenerationContext &genContext) {
1171   // Switch on op type being generated.
1172   if (opConfig.structuredOp) {
1173     return success(
1174         succeeded(generateNamedGenericOpOds(opConfig, genContext)) &&
1175         succeeded(generateNamedGenericOpDefns(opConfig, genContext)));
1176   }
1177   return emitError(genContext.getLoc()) << "unsupported operation type";
1178 }
1179
1180 //===----------------------------------------------------------------------===//
1181 // Command line options and main
1182 //===----------------------------------------------------------------------===//
1183
1184 static llvm::cl::opt<std::string>
1185     inputFilename(llvm::cl::Positional, llvm::cl::desc("<input file>"),
1186                   llvm::cl::init("-"), llvm::cl::value_desc("YAML filename"));
1187
1188 static llvm::cl::opt<std::string>
1189     outputOdsDeclFilename("o-ods-decl", llvm::cl::desc("ODS output filename"),
1190                           llvm::cl::value_desc("filename"), llvm::cl::init(""));
1191
1192 static llvm::cl::opt<std::string>
1193     outputCppImplFilename("o-impl",
1194                           llvm::cl::desc("C++ implementation file name"),
1195                           llvm::cl::value_desc("filename"), llvm::cl::init(""));
1196
1197 int main(int argc, char **argv) {
1198   llvm::cl::ParseCommandLineOptions(argc, argv, "Linalg ODS Gen from YAML");
1199
1200   // Set up the input file.
1201   std::string errorMessage;
1202   std::unique_ptr<llvm::MemoryBuffer> file =
1203       mlir::openInputFile(inputFilename, &errorMessage);
1204   if (!file) {
1205     llvm::errs() << errorMessage << "\n";
1206     return 1;
1207   }
1208
1209   MLIRContext mlirContext;
1210   LinalgYAMLContext yamlContext{&mlirContext};
1211
1212   std::vector<LinalgOpConfig> opConfigs;
1213
1214   // Parse input.
1215   Input yin(file->getBuffer(), &yamlContext);
1216   yin >> opConfigs;
1217
1218   if (yin.error())
1219     return 1;
1220
1221   // Open output files.
1222   std::unique_ptr<llvm::ToolOutputFile> outputOdsDecl;
1223   if (!outputOdsDeclFilename.empty()) {
1224     outputOdsDecl = openOutputFile(outputOdsDeclFilename, &errorMessage);
1225     if (!outputOdsDecl) {
1226       llvm::errs() << errorMessage << "\n";
1227       return 1;
1228     }
1229   }
1230
1231   std::unique_ptr<llvm::ToolOutputFile> outputCppImpl;
1232   if (!outputCppImplFilename.empty()) {
1233     outputCppImpl = openOutputFile(outputCppImplFilename, &errorMessage);
1234     if (!outputCppImpl) {
1235       llvm::errs() << errorMessage << "\n";
1236       return 1;
1237     }
1238   }
1239
1240   if (!outputOdsDecl && !outputCppImpl) {
1241     llvm::errs() << "error: No output files specified\n";
1242     return 1;
1243   }
1244
1245   // Generate.
1246   GenerationContext genContext(&mlirContext,
1247                                outputOdsDecl ? &outputOdsDecl->os() : nullptr,
1248                                outputCppImpl ? &outputCppImpl->os() : nullptr);
1249
1250   for (auto &opConfig : opConfigs) {
1251     if (!opConfig.metadata) {
1252       emitError(genContext.getLoc())
1253           << "missing operation metadata on subsequent op";
1254       return 1;
1255     }
1256
1257     genContext.setLoc(NameLoc::get(
1258         StringAttr::get(&mlirContext, opConfig.metadata->cppClassName)));
1259     if (failed(generateOp(opConfig, genContext))) {
1260       return 1;
1261     }
1262   }
1263
1264   if (outputOdsDecl)
1265     outputOdsDecl->keep();
1266   if (outputCppImpl)
1267     outputCppImpl->keep();
1268
1269   return 0;
1270 }