xref: /llvm-project/mlir/docs/Tutorials/Toy/Ch-7.md (revision bbfa7ef16dd9900b36abfa1a5f2faddb81afeb51)
15b4a01d4SMehdi Amini# Chapter 7: Adding a Composite Type to Toy
25b4a01d4SMehdi Amini
35b4a01d4SMehdi Amini[TOC]
45b4a01d4SMehdi Amini
55b4a01d4SMehdi AminiIn the [previous chapter](Ch-6.md), we demonstrated an end-to-end compilation
65b4a01d4SMehdi Aminiflow from our Toy front-end to LLVM IR. In this chapter, we will extend the Toy
75b4a01d4SMehdi Aminilanguage to support a new composite `struct` type.
85b4a01d4SMehdi Amini
95b4a01d4SMehdi Amini## Defining a `struct` in Toy
105b4a01d4SMehdi Amini
115b4a01d4SMehdi AminiThe first thing we need to define is the interface of this type in our `toy`
125b4a01d4SMehdi Aminisource language. The general syntax of a `struct` type in Toy is as follows:
135b4a01d4SMehdi Amini
145b4a01d4SMehdi Amini```toy
155b4a01d4SMehdi Amini# A struct is defined by using the `struct` keyword followed by a name.
165b4a01d4SMehdi Aministruct MyStruct {
175b4a01d4SMehdi Amini  # Inside of the struct is a list of variable declarations without initializers
185b4a01d4SMehdi Amini  # or shapes, which may also be other previously defined structs.
195b4a01d4SMehdi Amini  var a;
205b4a01d4SMehdi Amini  var b;
215b4a01d4SMehdi Amini}
225b4a01d4SMehdi Amini```
235b4a01d4SMehdi Amini
245b4a01d4SMehdi AminiStructs may now be used in functions as variables or parameters by using the
255b4a01d4SMehdi Amininame of the struct instead of `var`. The members of the struct are accessed via
265b4a01d4SMehdi Aminia `.` access operator. Values of `struct` type may be initialized with a
275b4a01d4SMehdi Aminicomposite initializer, or a comma-separated list of other initializers
285b4a01d4SMehdi Aminisurrounded by `{}`. An example is shown below:
295b4a01d4SMehdi Amini
305b4a01d4SMehdi Amini```toy
315b4a01d4SMehdi Aministruct Struct {
325b4a01d4SMehdi Amini  var a;
335b4a01d4SMehdi Amini  var b;
345b4a01d4SMehdi Amini}
355b4a01d4SMehdi Amini
365b4a01d4SMehdi Amini# User defined generic function may operate on struct types as well.
375b4a01d4SMehdi Aminidef multiply_transpose(Struct value) {
385b4a01d4SMehdi Amini  # We can access the elements of a struct via the '.' operator.
395b4a01d4SMehdi Amini  return transpose(value.a) * transpose(value.b);
405b4a01d4SMehdi Amini}
415b4a01d4SMehdi Amini
425b4a01d4SMehdi Aminidef main() {
435b4a01d4SMehdi Amini  # We initialize struct values using a composite initializer.
445b4a01d4SMehdi Amini  Struct value = {[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]};
455b4a01d4SMehdi Amini
465b4a01d4SMehdi Amini  # We pass these arguments to functions like we do with variables.
475b4a01d4SMehdi Amini  var c = multiply_transpose(value);
485b4a01d4SMehdi Amini  print(c);
495b4a01d4SMehdi Amini}
505b4a01d4SMehdi Amini```
515b4a01d4SMehdi Amini
525b4a01d4SMehdi Amini## Defining a `struct` in MLIR
535b4a01d4SMehdi Amini
545b4a01d4SMehdi AminiIn MLIR, we will also need a representation for our struct types. MLIR does not
555b4a01d4SMehdi Aminiprovide a type that does exactly what we need, so we will need to define our
565b4a01d4SMehdi Aminiown. We will simply define our `struct` as an unnamed container of a set of
575b4a01d4SMehdi Aminielement types. The name of the `struct` and its elements are only useful for the
585b4a01d4SMehdi AminiAST of our `toy` compiler, so we don't need to encode it in the MLIR
595b4a01d4SMehdi Aminirepresentation.
605b4a01d4SMehdi Amini
615b4a01d4SMehdi Amini### Defining the Type Class
625b4a01d4SMehdi Amini
635b4a01d4SMehdi Amini#### Defining the Type Class
645b4a01d4SMehdi Amini
6531d1ae79SMarkus BöckAs mentioned in [chapter 2](Ch-2.md), [`Type`](../../LangRef.md/#type-system)
665b4a01d4SMehdi Aminiobjects in MLIR are value-typed and rely on having an internal storage object
675b4a01d4SMehdi Aminithat holds the actual data for the type. The `Type` class in itself acts as a
685b4a01d4SMehdi Aminisimple wrapper around an internal `TypeStorage` object that is uniqued within an
695b4a01d4SMehdi Aminiinstance of an `MLIRContext`. When constructing a `Type`, we are internally just
705b4a01d4SMehdi Aminiconstructing and uniquing an instance of a storage class.
715b4a01d4SMehdi Amini
72c996d49cSRiver RiddleWhen defining a new `Type` that contains parametric data (e.g. the `struct`
73c996d49cSRiver Riddletype, which requires additional information to hold the element types), we will
74c996d49cSRiver Riddleneed to provide a derived storage class. The `singleton` types that don't have
7531d1ae79SMarkus Böckany additional data (e.g. the [`index` type](../../Dialects/Builtin.md/#indextype)) don't
76c996d49cSRiver Riddlerequire a storage class and use the default `TypeStorage`.
775b4a01d4SMehdi Amini
785b4a01d4SMehdi Amini##### Defining the Storage Class
795b4a01d4SMehdi Amini
805b4a01d4SMehdi AminiType storage objects contain all of the data necessary to construct and unique a
815b4a01d4SMehdi Aminitype instance. Derived storage classes must inherit from the base
825b4a01d4SMehdi Amini`mlir::TypeStorage` and provide a set of aliases and hooks that will be used by
835b4a01d4SMehdi Aminithe `MLIRContext` for uniquing. Below is the definition of the storage instance
845b4a01d4SMehdi Aminifor our `struct` type, with each of the necessary requirements detailed inline:
855b4a01d4SMehdi Amini
865b4a01d4SMehdi Amini```c++
875b4a01d4SMehdi Amini/// This class represents the internal storage of the Toy `StructType`.
885b4a01d4SMehdi Aministruct StructTypeStorage : public mlir::TypeStorage {
895b4a01d4SMehdi Amini  /// The `KeyTy` is a required type that provides an interface for the storage
905b4a01d4SMehdi Amini  /// instance. This type will be used when uniquing an instance of the type
915b4a01d4SMehdi Amini  /// storage. For our struct type, we will unique each instance structurally on
925b4a01d4SMehdi Amini  /// the elements that it contains.
935b4a01d4SMehdi Amini  using KeyTy = llvm::ArrayRef<mlir::Type>;
945b4a01d4SMehdi Amini
955b4a01d4SMehdi Amini  /// A constructor for the type storage instance.
965b4a01d4SMehdi Amini  StructTypeStorage(llvm::ArrayRef<mlir::Type> elementTypes)
975b4a01d4SMehdi Amini      : elementTypes(elementTypes) {}
985b4a01d4SMehdi Amini
995b4a01d4SMehdi Amini  /// Define the comparison function for the key type with the current storage
1005b4a01d4SMehdi Amini  /// instance. This is used when constructing a new instance to ensure that we
1015b4a01d4SMehdi Amini  /// haven't already uniqued an instance of the given key.
1025b4a01d4SMehdi Amini  bool operator==(const KeyTy &key) const { return key == elementTypes; }
1035b4a01d4SMehdi Amini
1045b4a01d4SMehdi Amini  /// Define a hash function for the key type. This is used when uniquing
1055b4a01d4SMehdi Amini  /// instances of the storage.
1065b4a01d4SMehdi Amini  /// Note: This method isn't necessary as both llvm::ArrayRef and mlir::Type
1075b4a01d4SMehdi Amini  /// have hash functions available, so we could just omit this entirely.
1085b4a01d4SMehdi Amini  static llvm::hash_code hashKey(const KeyTy &key) {
1095b4a01d4SMehdi Amini    return llvm::hash_value(key);
1105b4a01d4SMehdi Amini  }
1115b4a01d4SMehdi Amini
1125b4a01d4SMehdi Amini  /// Define a construction function for the key type from a set of parameters.
1135b4a01d4SMehdi Amini  /// These parameters will be provided when constructing the storage instance
1145b4a01d4SMehdi Amini  /// itself, see the `StructType::get` method further below.
1155b4a01d4SMehdi Amini  /// Note: This method isn't necessary because KeyTy can be directly
1165b4a01d4SMehdi Amini  /// constructed with the given parameters.
1175b4a01d4SMehdi Amini  static KeyTy getKey(llvm::ArrayRef<mlir::Type> elementTypes) {
1185b4a01d4SMehdi Amini    return KeyTy(elementTypes);
1195b4a01d4SMehdi Amini  }
1205b4a01d4SMehdi Amini
1215b4a01d4SMehdi Amini  /// Define a construction method for creating a new instance of this storage.
1225b4a01d4SMehdi Amini  /// This method takes an instance of a storage allocator, and an instance of a
1235b4a01d4SMehdi Amini  /// `KeyTy`. The given allocator must be used for *all* necessary dynamic
1245b4a01d4SMehdi Amini  /// allocations used to create the type storage and its internal.
1255b4a01d4SMehdi Amini  static StructTypeStorage *construct(mlir::TypeStorageAllocator &allocator,
1265b4a01d4SMehdi Amini                                      const KeyTy &key) {
1275b4a01d4SMehdi Amini    // Copy the elements from the provided `KeyTy` into the allocator.
1285b4a01d4SMehdi Amini    llvm::ArrayRef<mlir::Type> elementTypes = allocator.copyInto(key);
1295b4a01d4SMehdi Amini
1305b4a01d4SMehdi Amini    // Allocate the storage instance and construct it.
1315b4a01d4SMehdi Amini    return new (allocator.allocate<StructTypeStorage>())
1325b4a01d4SMehdi Amini        StructTypeStorage(elementTypes);
1335b4a01d4SMehdi Amini  }
1345b4a01d4SMehdi Amini
1355b4a01d4SMehdi Amini  /// The following field contains the element types of the struct.
1365b4a01d4SMehdi Amini  llvm::ArrayRef<mlir::Type> elementTypes;
1375b4a01d4SMehdi Amini};
1385b4a01d4SMehdi Amini```
1395b4a01d4SMehdi Amini
1405b4a01d4SMehdi Amini##### Defining the Type Class
1415b4a01d4SMehdi Amini
1425b4a01d4SMehdi AminiWith the storage class defined, we can add the definition for the user-visible
1435b4a01d4SMehdi Amini`StructType` class. This is the class that we will actually interface with.
1445b4a01d4SMehdi Amini
1455b4a01d4SMehdi Amini```c++
1465b4a01d4SMehdi Amini/// This class defines the Toy struct type. It represents a collection of
1475b4a01d4SMehdi Amini/// element types. All derived types in MLIR must inherit from the CRTP class
1485b4a01d4SMehdi Amini/// 'Type::TypeBase'. It takes as template parameters the concrete type
1495b4a01d4SMehdi Amini/// (StructType), the base class to use (Type), and the storage class
1505b4a01d4SMehdi Amini/// (StructTypeStorage).
1515b4a01d4SMehdi Aminiclass StructType : public mlir::Type::TypeBase<StructType, mlir::Type,
1525b4a01d4SMehdi Amini                                               StructTypeStorage> {
1535b4a01d4SMehdi Aminipublic:
1545b4a01d4SMehdi Amini  /// Inherit some necessary constructors from 'TypeBase'.
1555b4a01d4SMehdi Amini  using Base::Base;
1565b4a01d4SMehdi Amini
1575b4a01d4SMehdi Amini  /// Create an instance of a `StructType` with the given element types. There
1585b4a01d4SMehdi Amini  /// *must* be at least one element type.
1595b4a01d4SMehdi Amini  static StructType get(llvm::ArrayRef<mlir::Type> elementTypes) {
1605b4a01d4SMehdi Amini    assert(!elementTypes.empty() && "expected at least 1 element type");
1615b4a01d4SMehdi Amini
1625b4a01d4SMehdi Amini    // Call into a helper 'get' method in 'TypeBase' to get a uniqued instance
163250f43d3SRiver Riddle    // of this type. The first parameter is the context to unique in. The
164c996d49cSRiver Riddle    // parameters after are forwarded to the storage instance.
1655b4a01d4SMehdi Amini    mlir::MLIRContext *ctx = elementTypes.front().getContext();
166250f43d3SRiver Riddle    return Base::get(ctx, elementTypes);
1675b4a01d4SMehdi Amini  }
1685b4a01d4SMehdi Amini
1695b4a01d4SMehdi Amini  /// Returns the element types of this struct type.
1705b4a01d4SMehdi Amini  llvm::ArrayRef<mlir::Type> getElementTypes() {
1715b4a01d4SMehdi Amini    // 'getImpl' returns a pointer to the internal storage instance.
1725b4a01d4SMehdi Amini    return getImpl()->elementTypes;
1735b4a01d4SMehdi Amini  }
1745b4a01d4SMehdi Amini
1755b4a01d4SMehdi Amini  /// Returns the number of element type held by this struct.
1765b4a01d4SMehdi Amini  size_t getNumElementTypes() { return getElementTypes().size(); }
1775b4a01d4SMehdi Amini};
1785b4a01d4SMehdi Amini```
1795b4a01d4SMehdi Amini
180ee748605SRiver RiddleWe register this type in the `ToyDialect` initializer in a similar way to how we
1815b4a01d4SMehdi Aminidid with operations:
1825b4a01d4SMehdi Amini
1835b4a01d4SMehdi Amini```c++
184ee748605SRiver Riddlevoid ToyDialect::initialize() {
1855b4a01d4SMehdi Amini  addTypes<StructType>();
1865b4a01d4SMehdi Amini}
1875b4a01d4SMehdi Amini```
1885b4a01d4SMehdi Amini
18931bb8efdSRiver Riddle(An important note here is that when registering a type, the definition of the
19031bb8efdSRiver Riddlestorage class must be visible.)
19131bb8efdSRiver Riddle
1925b4a01d4SMehdi AminiWith this we can now use our `StructType` when generating MLIR from Toy. See
1935b4a01d4SMehdi Aminiexamples/toy/Ch7/mlir/MLIRGen.cpp for more details.
1945b4a01d4SMehdi Amini
195ee748605SRiver Riddle### Exposing to ODS
196ee748605SRiver Riddle
197ee748605SRiver RiddleAfter defining a new type, we should make the ODS framework aware of our Type so
198ee748605SRiver Riddlethat we can use it in the operation definitions and auto-generate utilities
199ee748605SRiver Riddlewithin the Dialect. A simple example is shown below:
200ee748605SRiver Riddle
201ee748605SRiver Riddle```tablegen
202ee748605SRiver Riddle// Provide a definition for the Toy StructType for use in ODS. This allows for
203ee748605SRiver Riddle// using StructType in a similar way to Tensor or MemRef. We use `DialectType`
204ee748605SRiver Riddle// to demarcate the StructType as belonging to the Toy dialect.
205ee748605SRiver Riddledef Toy_StructType :
206ee748605SRiver Riddle    DialectType<Toy_Dialect, CPred<"$_self.isa<StructType>()">,
207ee748605SRiver Riddle                "Toy struct type">;
208ee748605SRiver Riddle
209ee748605SRiver Riddle// Provide a definition of the types that are used within the Toy dialect.
210ee748605SRiver Riddledef Toy_Type : AnyTypeOf<[F64Tensor, Toy_StructType]>;
211ee748605SRiver Riddle```
212ee748605SRiver Riddle
2135b4a01d4SMehdi Amini### Parsing and Printing
2145b4a01d4SMehdi Amini
2155b4a01d4SMehdi AminiAt this point we can use our `StructType` during MLIR generation and
2165b4a01d4SMehdi Aminitransformation, but we can't output or parse `.mlir`. For this we need to add
2175b4a01d4SMehdi Aminisupport for parsing and printing instances of the `StructType`. This can be done
2185b4a01d4SMehdi Aminiby overriding the `parseType` and `printType` methods on the `ToyDialect`.
219ee748605SRiver RiddleDeclarations for these methods are automatically provided when the type is
220ee748605SRiver Riddleexposed to ODS as detailed in the previous section.
2215b4a01d4SMehdi Amini
2225b4a01d4SMehdi Amini```c++
2235b4a01d4SMehdi Aminiclass ToyDialect : public mlir::Dialect {
2245b4a01d4SMehdi Aminipublic:
2255b4a01d4SMehdi Amini  /// Parse an instance of a type registered to the toy dialect.
2265b4a01d4SMehdi Amini  mlir::Type parseType(mlir::DialectAsmParser &parser) const override;
2275b4a01d4SMehdi Amini
2285b4a01d4SMehdi Amini  /// Print an instance of a type registered to the toy dialect.
2295b4a01d4SMehdi Amini  void printType(mlir::Type type,
2305b4a01d4SMehdi Amini                 mlir::DialectAsmPrinter &printer) const override;
2315b4a01d4SMehdi Amini};
2325b4a01d4SMehdi Amini```
2335b4a01d4SMehdi Amini
2345b4a01d4SMehdi AminiThese methods take an instance of a high-level parser or printer that allows for
2355b4a01d4SMehdi Aminieasily implementing the necessary functionality. Before going into the
2365b4a01d4SMehdi Aminiimplementation, let's think about the syntax that we want for the `struct` type
2375b4a01d4SMehdi Aminiin the printed IR. As described in the
23831d1ae79SMarkus Böck[MLIR language reference](../../LangRef.md/#dialect-types), dialect types are
2395b4a01d4SMehdi Aminigenerally represented as: `! dialect-namespace < type-data >`, with a pretty
2405b4a01d4SMehdi Aminiform available under certain circumstances. The responsibility of our `Toy`
2415b4a01d4SMehdi Aminiparser and printer is to provide the `type-data` bits. We will define our
2425b4a01d4SMehdi Amini`StructType` as having the following form:
2435b4a01d4SMehdi Amini
2445b4a01d4SMehdi Amini```
2455b4a01d4SMehdi Amini  struct-type ::= `struct` `<` type (`,` type)* `>`
2465b4a01d4SMehdi Amini```
2475b4a01d4SMehdi Amini
2485b4a01d4SMehdi Amini#### Parsing
2495b4a01d4SMehdi Amini
2505b4a01d4SMehdi AminiAn implementation of the parser is shown below:
2515b4a01d4SMehdi Amini
2525b4a01d4SMehdi Amini```c++
2535b4a01d4SMehdi Amini/// Parse an instance of a type registered to the toy dialect.
2545b4a01d4SMehdi Aminimlir::Type ToyDialect::parseType(mlir::DialectAsmParser &parser) const {
2555b4a01d4SMehdi Amini  // Parse a struct type in the following form:
2565b4a01d4SMehdi Amini  //   struct-type ::= `struct` `<` type (`,` type)* `>`
2575b4a01d4SMehdi Amini
2585b4a01d4SMehdi Amini  // NOTE: All MLIR parser function return a ParseResult. This is a
2595b4a01d4SMehdi Amini  // specialization of LogicalResult that auto-converts to a `true` boolean
2605b4a01d4SMehdi Amini  // value on failure to allow for chaining, but may be used with explicit
2615b4a01d4SMehdi Amini  // `mlir::failed/mlir::succeeded` as desired.
2625b4a01d4SMehdi Amini
2635b4a01d4SMehdi Amini  // Parse: `struct` `<`
2645b4a01d4SMehdi Amini  if (parser.parseKeyword("struct") || parser.parseLess())
2655b4a01d4SMehdi Amini    return Type();
2665b4a01d4SMehdi Amini
2675b4a01d4SMehdi Amini  // Parse the element types of the struct.
2685b4a01d4SMehdi Amini  SmallVector<mlir::Type, 1> elementTypes;
2695b4a01d4SMehdi Amini  do {
2705b4a01d4SMehdi Amini    // Parse the current element type.
2716842ec42SRiver Riddle    SMLoc typeLoc = parser.getCurrentLocation();
2725b4a01d4SMehdi Amini    mlir::Type elementType;
2735b4a01d4SMehdi Amini    if (parser.parseType(elementType))
2745b4a01d4SMehdi Amini      return nullptr;
2755b4a01d4SMehdi Amini
2765b4a01d4SMehdi Amini    // Check that the type is either a TensorType or another StructType.
277ee394e68SRahul Joshi    if (!elementType.isa<mlir::TensorType, StructType>()) {
2785b4a01d4SMehdi Amini      parser.emitError(typeLoc, "element type for a struct must either "
2795b4a01d4SMehdi Amini                                "be a TensorType or a StructType, got: ")
2805b4a01d4SMehdi Amini          << elementType;
2815b4a01d4SMehdi Amini      return Type();
2825b4a01d4SMehdi Amini    }
2835b4a01d4SMehdi Amini    elementTypes.push_back(elementType);
2845b4a01d4SMehdi Amini
2855b4a01d4SMehdi Amini    // Parse the optional: `,`
2865b4a01d4SMehdi Amini  } while (succeeded(parser.parseOptionalComma()));
2875b4a01d4SMehdi Amini
2885b4a01d4SMehdi Amini  // Parse: `>`
2895b4a01d4SMehdi Amini  if (parser.parseGreater())
2905b4a01d4SMehdi Amini    return Type();
2915b4a01d4SMehdi Amini  return StructType::get(elementTypes);
2925b4a01d4SMehdi Amini}
2935b4a01d4SMehdi Amini```
2945b4a01d4SMehdi Amini
2955b4a01d4SMehdi Amini#### Printing
2965b4a01d4SMehdi Amini
2975b4a01d4SMehdi AminiAn implementation of the printer is shown below:
2985b4a01d4SMehdi Amini
2995b4a01d4SMehdi Amini```c++
3005b4a01d4SMehdi Amini/// Print an instance of a type registered to the toy dialect.
3015b4a01d4SMehdi Aminivoid ToyDialect::printType(mlir::Type type,
3025b4a01d4SMehdi Amini                           mlir::DialectAsmPrinter &printer) const {
3035b4a01d4SMehdi Amini  // Currently the only toy type is a struct type.
3045b4a01d4SMehdi Amini  StructType structType = type.cast<StructType>();
3055b4a01d4SMehdi Amini
3065b4a01d4SMehdi Amini  // Print the struct type according to the parser format.
3075b4a01d4SMehdi Amini  printer << "struct<";
3082f21a579SRiver Riddle  llvm::interleaveComma(structType.getElementTypes(), printer);
3095b4a01d4SMehdi Amini  printer << '>';
3105b4a01d4SMehdi Amini}
3115b4a01d4SMehdi Amini```
3125b4a01d4SMehdi Amini
3135b4a01d4SMehdi AminiBefore moving on, let's look at a quick of example showcasing the functionality
3145b4a01d4SMehdi Aminiwe have now:
3155b4a01d4SMehdi Amini
3165b4a01d4SMehdi Amini```toy
3175b4a01d4SMehdi Aministruct Struct {
3185b4a01d4SMehdi Amini  var a;
3195b4a01d4SMehdi Amini  var b;
3205b4a01d4SMehdi Amini}
3215b4a01d4SMehdi Amini
3225b4a01d4SMehdi Aminidef multiply_transpose(Struct value) {
3235b4a01d4SMehdi Amini}
3245b4a01d4SMehdi Amini```
3255b4a01d4SMehdi Amini
3265b4a01d4SMehdi AminiWhich generates the following:
3275b4a01d4SMehdi Amini
3285b4a01d4SMehdi Amini```mlir
3295b4a01d4SMehdi Aminimodule {
330ee2c6cd9SRiver Riddle  toy.func @multiply_transpose(%arg0: !toy.struct<tensor<*xf64>, tensor<*xf64>>) {
3310050e8f0SRiver Riddle    toy.return
3325b4a01d4SMehdi Amini  }
3335b4a01d4SMehdi Amini}
3345b4a01d4SMehdi Amini```
3355b4a01d4SMehdi Amini
3365b4a01d4SMehdi Amini### Operating on `StructType`
3375b4a01d4SMehdi Amini
3385b4a01d4SMehdi AminiNow that the `struct` type has been defined, and we can round-trip it through
3395b4a01d4SMehdi Aminithe IR. The next step is to add support for using it within our operations.
3405b4a01d4SMehdi Amini
3415b4a01d4SMehdi Amini#### Updating Existing Operations
3425b4a01d4SMehdi Amini
343ee748605SRiver RiddleA few of our existing operations, e.g. `ReturnOp`, will need to be updated to
344ee748605SRiver Riddlehandle `Toy_StructType`.
3455b4a01d4SMehdi Amini
3465b4a01d4SMehdi Amini```tablegen
3475b4a01d4SMehdi Aminidef ReturnOp : Toy_Op<"return", [Terminator, HasParent<"FuncOp">]> {
3485b4a01d4SMehdi Amini  ...
3495b4a01d4SMehdi Amini  let arguments = (ins Variadic<Toy_Type>:$input);
3505b4a01d4SMehdi Amini  ...
3515b4a01d4SMehdi Amini}
3525b4a01d4SMehdi Amini```
3535b4a01d4SMehdi Amini
3545b4a01d4SMehdi Amini#### Adding New `Toy` Operations
3555b4a01d4SMehdi Amini
3565b4a01d4SMehdi AminiIn addition to the existing operations, we will be adding a few new operations
3575b4a01d4SMehdi Aminithat will provide more specific handling of `structs`.
3585b4a01d4SMehdi Amini
3595b4a01d4SMehdi Amini##### `toy.struct_constant`
3605b4a01d4SMehdi Amini
3615b4a01d4SMehdi AminiThis new operation materializes a constant value for a struct. In our current
36231d1ae79SMarkus Böckmodeling, we just use an [array attribute](../../Dialects/Builtin.md/#arrayattr)
3635b4a01d4SMehdi Aminithat contains a set of constant values for each of the `struct` elements.
3645b4a01d4SMehdi Amini
3655b4a01d4SMehdi Amini```mlir
3660050e8f0SRiver Riddle  %0 = toy.struct_constant [
3670050e8f0SRiver Riddle    dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf64>
3680050e8f0SRiver Riddle  ] : !toy.struct<tensor<*xf64>>
3695b4a01d4SMehdi Amini```
3705b4a01d4SMehdi Amini
3715b4a01d4SMehdi Amini##### `toy.struct_access`
3725b4a01d4SMehdi Amini
3735b4a01d4SMehdi AminiThis new operation materializes the Nth element of a `struct` value.
3745b4a01d4SMehdi Amini
3755b4a01d4SMehdi Amini```mlir
37645d522d6SMatthias Kramm  // Using %0 from above
3770050e8f0SRiver Riddle  %1 = toy.struct_access %0[0] : !toy.struct<tensor<*xf64>> -> tensor<*xf64>
3785b4a01d4SMehdi Amini```
3795b4a01d4SMehdi Amini
3805b4a01d4SMehdi AminiWith these operations, we can revisit our original example:
3815b4a01d4SMehdi Amini
3825b4a01d4SMehdi Amini```toy
3835b4a01d4SMehdi Aministruct Struct {
3845b4a01d4SMehdi Amini  var a;
3855b4a01d4SMehdi Amini  var b;
3865b4a01d4SMehdi Amini}
3875b4a01d4SMehdi Amini
3885b4a01d4SMehdi Amini# User defined generic function may operate on struct types as well.
3895b4a01d4SMehdi Aminidef multiply_transpose(Struct value) {
3905b4a01d4SMehdi Amini  # We can access the elements of a struct via the '.' operator.
3915b4a01d4SMehdi Amini  return transpose(value.a) * transpose(value.b);
3925b4a01d4SMehdi Amini}
3935b4a01d4SMehdi Amini
3945b4a01d4SMehdi Aminidef main() {
3955b4a01d4SMehdi Amini  # We initialize struct values using a composite initializer.
3965b4a01d4SMehdi Amini  Struct value = {[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]};
3975b4a01d4SMehdi Amini
3985b4a01d4SMehdi Amini  # We pass these arguments to functions like we do with variables.
3995b4a01d4SMehdi Amini  var c = multiply_transpose(value);
4005b4a01d4SMehdi Amini  print(c);
4015b4a01d4SMehdi Amini}
4025b4a01d4SMehdi Amini```
4035b4a01d4SMehdi Amini
4045b4a01d4SMehdi Aminiand finally get a full MLIR module:
4055b4a01d4SMehdi Amini
4065b4a01d4SMehdi Amini```mlir
4075b4a01d4SMehdi Aminimodule {
408ee2c6cd9SRiver Riddle  toy.func @multiply_transpose(%arg0: !toy.struct<tensor<*xf64>, tensor<*xf64>>) -> tensor<*xf64> {
4090050e8f0SRiver Riddle    %0 = toy.struct_access %arg0[0] : !toy.struct<tensor<*xf64>, tensor<*xf64>> -> tensor<*xf64>
4100050e8f0SRiver Riddle    %1 = toy.transpose(%0 : tensor<*xf64>) to tensor<*xf64>
4110050e8f0SRiver Riddle    %2 = toy.struct_access %arg0[1] : !toy.struct<tensor<*xf64>, tensor<*xf64>> -> tensor<*xf64>
4120050e8f0SRiver Riddle    %3 = toy.transpose(%2 : tensor<*xf64>) to tensor<*xf64>
4130050e8f0SRiver Riddle    %4 = toy.mul %1, %3 : tensor<*xf64>
4140050e8f0SRiver Riddle    toy.return %4 : tensor<*xf64>
4155b4a01d4SMehdi Amini  }
416ee2c6cd9SRiver Riddle  toy.func @main() {
4170050e8f0SRiver Riddle    %0 = toy.struct_constant [
4180050e8f0SRiver Riddle      dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>,
4190050e8f0SRiver Riddle      dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
4200050e8f0SRiver Riddle    ] : !toy.struct<tensor<*xf64>, tensor<*xf64>>
4210050e8f0SRiver Riddle    %1 = toy.generic_call @multiply_transpose(%0) : (!toy.struct<tensor<*xf64>, tensor<*xf64>>) -> tensor<*xf64>
4220050e8f0SRiver Riddle    toy.print %1 : tensor<*xf64>
4230050e8f0SRiver Riddle    toy.return
4245b4a01d4SMehdi Amini  }
4255b4a01d4SMehdi Amini}
4265b4a01d4SMehdi Amini```
4275b4a01d4SMehdi Amini
4285b4a01d4SMehdi Amini#### Optimizing Operations on `StructType`
4295b4a01d4SMehdi Amini
4305b4a01d4SMehdi AminiNow that we have a few operations operating on `StructType`, we also have many
4315b4a01d4SMehdi Amininew constant folding opportunities.
4325b4a01d4SMehdi Amini
4335b4a01d4SMehdi AminiAfter inlining, the MLIR module in the previous section looks something like:
4345b4a01d4SMehdi Amini
4355b4a01d4SMehdi Amini```mlir
4365b4a01d4SMehdi Aminimodule {
437ee2c6cd9SRiver Riddle  toy.func @main() {
4380050e8f0SRiver Riddle    %0 = toy.struct_constant [
4390050e8f0SRiver Riddle      dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>,
4400050e8f0SRiver Riddle      dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
4410050e8f0SRiver Riddle    ] : !toy.struct<tensor<*xf64>, tensor<*xf64>>
4420050e8f0SRiver Riddle    %1 = toy.struct_access %0[0] : !toy.struct<tensor<*xf64>, tensor<*xf64>> -> tensor<*xf64>
4430050e8f0SRiver Riddle    %2 = toy.transpose(%1 : tensor<*xf64>) to tensor<*xf64>
4440050e8f0SRiver Riddle    %3 = toy.struct_access %0[1] : !toy.struct<tensor<*xf64>, tensor<*xf64>> -> tensor<*xf64>
4450050e8f0SRiver Riddle    %4 = toy.transpose(%3 : tensor<*xf64>) to tensor<*xf64>
4460050e8f0SRiver Riddle    %5 = toy.mul %2, %4 : tensor<*xf64>
4470050e8f0SRiver Riddle    toy.print %5 : tensor<*xf64>
4480050e8f0SRiver Riddle    toy.return
4495b4a01d4SMehdi Amini  }
4505b4a01d4SMehdi Amini}
4515b4a01d4SMehdi Amini```
4525b4a01d4SMehdi Amini
4535b4a01d4SMehdi AminiWe have several `toy.struct_access` operations that access into a
45445d522d6SMatthias Kramm`toy.struct_constant`. As detailed in [chapter 3](Ch-3.md) (FoldConstantReshape),
45545d522d6SMatthias Krammwe can add folders for these `toy` operations by setting the `hasFolder` bit
45645d522d6SMatthias Krammon the operation definition and providing a definition of the `*Op::fold`
45745d522d6SMatthias Krammmethod.
4585b4a01d4SMehdi Amini
4595b4a01d4SMehdi Amini```c++
4605b4a01d4SMehdi Amini/// Fold constants.
461*bbfa7ef1SMarkus BöckOpFoldResult ConstantOp::fold(FoldAdaptor adaptor) { return value(); }
4625b4a01d4SMehdi Amini
4635b4a01d4SMehdi Amini/// Fold struct constants.
464*bbfa7ef1SMarkus BöckOpFoldResult StructConstantOp::fold(FoldAdaptor adaptor) {
4655b4a01d4SMehdi Amini  return value();
4665b4a01d4SMehdi Amini}
4675b4a01d4SMehdi Amini
4685b4a01d4SMehdi Amini/// Fold simple struct access operations that access into a constant.
469*bbfa7ef1SMarkus BöckOpFoldResult StructAccessOp::fold(FoldAdaptor adaptor) {
470*bbfa7ef1SMarkus Böck  auto structAttr = adaptor.getInput().dyn_cast_or_null<mlir::ArrayAttr>();
4715b4a01d4SMehdi Amini  if (!structAttr)
4725b4a01d4SMehdi Amini    return nullptr;
4735b4a01d4SMehdi Amini
4745b4a01d4SMehdi Amini  size_t elementIndex = index().getZExtValue();
4752101590aSUday Bondhugula  return structAttr[elementIndex];
4765b4a01d4SMehdi Amini}
4775b4a01d4SMehdi Amini```
4785b4a01d4SMehdi Amini
4795b4a01d4SMehdi AminiTo ensure that MLIR generates the proper constant operations when folding our
4805b4a01d4SMehdi Amini`Toy` operations, i.e. `ConstantOp` for `TensorType` and `StructConstant` for
4815b4a01d4SMehdi Amini`StructType`, we will need to provide an override for the dialect hook
4825b4a01d4SMehdi Amini`materializeConstant`. This allows for generic MLIR operations to create
4835b4a01d4SMehdi Aminiconstants for the `Toy` dialect when necessary.
4845b4a01d4SMehdi Amini
4855b4a01d4SMehdi Amini```c++
4865b4a01d4SMehdi Aminimlir::Operation *ToyDialect::materializeConstant(mlir::OpBuilder &builder,
4875b4a01d4SMehdi Amini                                                 mlir::Attribute value,
4885b4a01d4SMehdi Amini                                                 mlir::Type type,
4895b4a01d4SMehdi Amini                                                 mlir::Location loc) {
4905b4a01d4SMehdi Amini  if (type.isa<StructType>())
4915b4a01d4SMehdi Amini    return builder.create<StructConstantOp>(loc, type,
4925b4a01d4SMehdi Amini                                            value.cast<mlir::ArrayAttr>());
4935b4a01d4SMehdi Amini  return builder.create<ConstantOp>(loc, type,
4945b4a01d4SMehdi Amini                                    value.cast<mlir::DenseElementsAttr>());
4955b4a01d4SMehdi Amini}
4965b4a01d4SMehdi Amini```
4975b4a01d4SMehdi Amini
4985b4a01d4SMehdi AminiWith this, we can now generate code that can be generated to LLVM without any
4995b4a01d4SMehdi Aminichanges to our pipeline.
5005b4a01d4SMehdi Amini
5015b4a01d4SMehdi Amini```mlir
5025b4a01d4SMehdi Aminimodule {
503ee2c6cd9SRiver Riddle  toy.func @main() {
5040050e8f0SRiver Riddle    %0 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
5050050e8f0SRiver Riddle    %1 = toy.transpose(%0 : tensor<2x3xf64>) to tensor<3x2xf64>
5060050e8f0SRiver Riddle    %2 = toy.mul %1, %1 : tensor<3x2xf64>
5070050e8f0SRiver Riddle    toy.print %2 : tensor<3x2xf64>
5080050e8f0SRiver Riddle    toy.return
5095b4a01d4SMehdi Amini  }
5105b4a01d4SMehdi Amini}
5115b4a01d4SMehdi Amini```
5125b4a01d4SMehdi Amini
5135b4a01d4SMehdi AminiYou can build `toyc-ch7` and try yourself: `toyc-ch7
5145b4a01d4SMehdi Aminitest/Examples/Toy/Ch7/struct-codegen.toy -emit=mlir`. More details on defining
5155b4a01d4SMehdi Aminicustom types can be found in
5161294fa69SRiver Riddle[DefiningAttributesAndTypes](../../DefiningDialects/AttributesAndTypes.md).
517