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