517 lines
19 KiB
Markdown
517 lines
19 KiB
Markdown
|
# Chapter 7: Adding a Composite Type to Toy
|
||
|
|
|
||
|
|
[TOC]
|
||
|
|
|
||
|
|
In the [previous chapter](Ch-6.md), we demonstrated an end-to-end compilation
|
||
|
|
flow from our Toy front-end to LLVM IR. In this chapter, we will extend the Toy
|
||
|
|
language to support a new composite `struct` type.
|
||
|
|
|
||
|
|
## Defining a `struct` in Toy
|
||
|
|
|
||
|
|
The first thing we need to define is the interface of this type in our `toy`
|
||
|
|
source language. The general syntax of a `struct` type in Toy is as follows:
|
||
|
|
|
||
|
|
```toy
|
||
|
|
# A struct is defined by using the `struct` keyword followed by a name.
|
||
|
|
struct MyStruct {
|
||
|
|
# Inside of the struct is a list of variable declarations without initializers
|
||
|
|
# or shapes, which may also be other previously defined structs.
|
||
|
|
var a;
|
||
|
|
var b;
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
Structs may now be used in functions as variables or parameters by using the
|
||
|
|
name of the struct instead of `var`. The members of the struct are accessed via
|
||
|
|
a `.` access operator. Values of `struct` type may be initialized with a
|
||
|
|
composite initializer, or a comma-separated list of other initializers
|
||
|
|
surrounded by `{}`. An example is shown below:
|
||
|
|
|
||
|
|
```toy
|
||
|
|
struct Struct {
|
||
|
|
var a;
|
||
|
|
var b;
|
||
|
|
}
|
||
|
|
|
||
|
|
# User defined generic function may operate on struct types as well.
|
||
|
|
def multiply_transpose(Struct value) {
|
||
|
|
# We can access the elements of a struct via the '.' operator.
|
||
|
|
return transpose(value.a) * transpose(value.b);
|
||
|
|
}
|
||
|
|
|
||
|
|
def main() {
|
||
|
|
# We initialize struct values using a composite initializer.
|
||
|
|
Struct value = {[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]};
|
||
|
|
|
||
|
|
# We pass these arguments to functions like we do with variables.
|
||
|
|
var c = multiply_transpose(value);
|
||
|
|
print(c);
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
## Defining a `struct` in MLIR
|
||
|
|
|
||
|
|
In MLIR, we will also need a representation for our struct types. MLIR does not
|
||
|
|
provide a type that does exactly what we need, so we will need to define our
|
||
|
|
own. We will simply define our `struct` as an unnamed container of a set of
|
||
|
|
element types. The name of the `struct` and its elements are only useful for the
|
||
|
|
AST of our `toy` compiler, so we don't need to encode it in the MLIR
|
||
|
|
representation.
|
||
|
|
|
||
|
|
### Defining the Type Class
|
||
|
|
|
||
|
|
#### Defining the Type Class
|
||
|
|
|
||
|
|
As mentioned in [chapter 2](Ch-2.md), [`Type`](../../LangRef.md/#type-system)
|
||
|
|
objects in MLIR are value-typed and rely on having an internal storage object
|
||
|
|
that holds the actual data for the type. The `Type` class in itself acts as a
|
||
|
|
simple wrapper around an internal `TypeStorage` object that is uniqued within an
|
||
|
|
instance of an `MLIRContext`. When constructing a `Type`, we are internally just
|
||
|
|
constructing and uniquing an instance of a storage class.
|
||
|
|
|
||
|
|
When defining a new `Type` that contains parametric data (e.g. the `struct`
|
||
|
|
type, which requires additional information to hold the element types), we will
|
||
|
|
need to provide a derived storage class. The `singleton` types that don't have
|
||
|
|
any additional data (e.g. the [`index` type](../../Dialects/Builtin.md/#indextype)) don't
|
||
|
|
require a storage class and use the default `TypeStorage`.
|
||
|
|
|
||
|
|
##### Defining the Storage Class
|
||
|
|
|
||
|
|
Type storage objects contain all of the data necessary to construct and unique a
|
||
|
|
type instance. Derived storage classes must inherit from the base
|
||
|
|
`mlir::TypeStorage` and provide a set of aliases and hooks that will be used by
|
||
|
|
the `MLIRContext` for uniquing. Below is the definition of the storage instance
|
||
|
|
for our `struct` type, with each of the necessary requirements detailed inline:
|
||
|
|
|
||
|
|
```c++
|
||
|
|
/// This class represents the internal storage of the Toy `StructType`.
|
||
|
|
struct StructTypeStorage : public mlir::TypeStorage {
|
||
|
|
/// The `KeyTy` is a required type that provides an interface for the storage
|
||
|
|
/// instance. This type will be used when uniquing an instance of the type
|
||
|
|
/// storage. For our struct type, we will unique each instance structurally on
|
||
|
|
/// the elements that it contains.
|
||
|
|
using KeyTy = llvm::ArrayRef<mlir::Type>;
|
||
|
|
|
||
|
|
/// A constructor for the type storage instance.
|
||
|
|
StructTypeStorage(llvm::ArrayRef<mlir::Type> elementTypes)
|
||
|
|
: elementTypes(elementTypes) {}
|
||
|
|
|
||
|
|
/// Define the comparison function for the key type with the current storage
|
||
|
|
/// instance. This is used when constructing a new instance to ensure that we
|
||
|
|
/// haven't already uniqued an instance of the given key.
|
||
|
|
bool operator==(const KeyTy &key) const { return key == elementTypes; }
|
||
|
|
|
||
|
|
/// Define a hash function for the key type. This is used when uniquing
|
||
|
|
/// instances of the storage.
|
||
|
|
/// Note: This method isn't necessary as both llvm::ArrayRef and mlir::Type
|
||
|
|
/// have hash functions available, so we could just omit this entirely.
|
||
|
|
static llvm::hash_code hashKey(const KeyTy &key) {
|
||
|
|
return llvm::hash_value(key);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Define a construction function for the key type from a set of parameters.
|
||
|
|
/// These parameters will be provided when constructing the storage instance
|
||
|
|
/// itself, see the `StructType::get` method further below.
|
||
|
|
/// Note: This method isn't necessary because KeyTy can be directly
|
||
|
|
/// constructed with the given parameters.
|
||
|
|
static KeyTy getKey(llvm::ArrayRef<mlir::Type> elementTypes) {
|
||
|
|
return KeyTy(elementTypes);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Define a construction method for creating a new instance of this storage.
|
||
|
|
/// This method takes an instance of a storage allocator, and an instance of a
|
||
|
|
/// `KeyTy`. The given allocator must be used for *all* necessary dynamic
|
||
|
|
/// allocations used to create the type storage and its internal.
|
||
|
|
static StructTypeStorage *construct(mlir::TypeStorageAllocator &allocator,
|
||
|
|
const KeyTy &key) {
|
||
|
|
// Copy the elements from the provided `KeyTy` into the allocator.
|
||
|
|
llvm::ArrayRef<mlir::Type> elementTypes = allocator.copyInto(key);
|
||
|
|
|
||
|
|
// Allocate the storage instance and construct it.
|
||
|
|
return new (allocator.allocate<StructTypeStorage>())
|
||
|
|
StructTypeStorage(elementTypes);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// The following field contains the element types of the struct.
|
||
|
|
llvm::ArrayRef<mlir::Type> elementTypes;
|
||
|
|
};
|
||
|
|
```
|
||
|
|
|
||
|
|
##### Defining the Type Class
|
||
|
|
|
||
|
|
With the storage class defined, we can add the definition for the user-visible
|
||
|
|
`StructType` class. This is the class that we will actually interface with.
|
||
|
|
|
||
|
|
```c++
|
||
|
|
/// This class defines the Toy struct type. It represents a collection of
|
||
|
|
/// element types. All derived types in MLIR must inherit from the CRTP class
|
||
|
|
/// 'Type::TypeBase'. It takes as template parameters the concrete type
|
||
|
|
/// (StructType), the base class to use (Type), and the storage class
|
||
|
|
/// (StructTypeStorage).
|
||
|
|
class StructType : public mlir::Type::TypeBase<StructType, mlir::Type,
|
||
|
|
StructTypeStorage> {
|
||
|
|
public:
|
||
|
|
/// Inherit some necessary constructors from 'TypeBase'.
|
||
|
|
using Base::Base;
|
||
|
|
|
||
|
|
/// Create an instance of a `StructType` with the given element types. There
|
||
|
|
/// *must* be at least one element type.
|
||
|
|
static StructType get(llvm::ArrayRef<mlir::Type> elementTypes) {
|
||
|
|
assert(!elementTypes.empty() && "expected at least 1 element type");
|
||
|
|
|
||
|
|
// Call into a helper 'get' method in 'TypeBase' to get a uniqued instance
|
||
|
|
// of this type. The first parameter is the context to unique in. The
|
||
|
|
// parameters after are forwarded to the storage instance.
|
||
|
|
mlir::MLIRContext *ctx = elementTypes.front().getContext();
|
||
|
|
return Base::get(ctx, elementTypes);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the element types of this struct type.
|
||
|
|
llvm::ArrayRef<mlir::Type> getElementTypes() {
|
||
|
|
// 'getImpl' returns a pointer to the internal storage instance.
|
||
|
|
return getImpl()->elementTypes;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the number of element type held by this struct.
|
||
|
|
size_t getNumElementTypes() { return getElementTypes().size(); }
|
||
|
|
};
|
||
|
|
```
|
||
|
|
|
||
|
|
We register this type in the `ToyDialect` initializer in a similar way to how we
|
||
|
|
did with operations:
|
||
|
|
|
||
|
|
```c++
|
||
|
|
void ToyDialect::initialize() {
|
||
|
|
addTypes<StructType>();
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
(An important note here is that when registering a type, the definition of the
|
||
|
|
storage class must be visible.)
|
||
|
|
|
||
|
|
With this we can now use our `StructType` when generating MLIR from Toy. See
|
||
|
|
examples/toy/Ch7/mlir/MLIRGen.cpp for more details.
|
||
|
|
|
||
|
|
### Exposing to ODS
|
||
|
|
|
||
|
|
After defining a new type, we should make the ODS framework aware of our Type so
|
||
|
|
that we can use it in the operation definitions and auto-generate utilities
|
||
|
|
within the Dialect. A simple example is shown below:
|
||
|
|
|
||
|
|
```tablegen
|
||
|
|
// Provide a definition for the Toy StructType for use in ODS. This allows for
|
||
|
|
// using StructType in a similar way to Tensor or MemRef. We use `DialectType`
|
||
|
|
// to demarcate the StructType as belonging to the Toy dialect.
|
||
|
|
def Toy_StructType :
|
||
|
|
DialectType<Toy_Dialect, CPred<"$_self.isa<StructType>()">,
|
||
|
|
"Toy struct type">;
|
||
|
|
|
||
|
|
// Provide a definition of the types that are used within the Toy dialect.
|
||
|
|
def Toy_Type : AnyTypeOf<[F64Tensor, Toy_StructType]>;
|
||
|
|
```
|
||
|
|
|
||
|
|
### Parsing and Printing
|
||
|
|
|
||
|
|
At this point we can use our `StructType` during MLIR generation and
|
||
|
|
transformation, but we can't output or parse `.mlir`. For this we need to add
|
||
|
|
support for parsing and printing instances of the `StructType`. This can be done
|
||
|
|
by overriding the `parseType` and `printType` methods on the `ToyDialect`.
|
||
|
|
Declarations for these methods are automatically provided when the type is
|
||
|
|
exposed to ODS as detailed in the previous section.
|
||
|
|
|
||
|
|
```c++
|
||
|
|
class ToyDialect : public mlir::Dialect {
|
||
|
|
public:
|
||
|
|
/// Parse an instance of a type registered to the toy dialect.
|
||
|
|
mlir::Type parseType(mlir::DialectAsmParser &parser) const override;
|
||
|
|
|
||
|
|
/// Print an instance of a type registered to the toy dialect.
|
||
|
|
void printType(mlir::Type type,
|
||
|
|
mlir::DialectAsmPrinter &printer) const override;
|
||
|
|
};
|
||
|
|
```
|
||
|
|
|
||
|
|
These methods take an instance of a high-level parser or printer that allows for
|
||
|
|
easily implementing the necessary functionality. Before going into the
|
||
|
|
implementation, let's think about the syntax that we want for the `struct` type
|
||
|
|
in the printed IR. As described in the
|
||
|
|
[MLIR language reference](../../LangRef.md/#dialect-types), dialect types are
|
||
|
|
generally represented as: `! dialect-namespace < type-data >`, with a pretty
|
||
|
|
form available under certain circumstances. The responsibility of our `Toy`
|
||
|
|
parser and printer is to provide the `type-data` bits. We will define our
|
||
|
|
`StructType` as having the following form:
|
||
|
|
|
||
|
|
```
|
||
|
|
struct-type ::= `struct` `<` type (`,` type)* `>`
|
||
|
|
```
|
||
|
|
|
||
|
|
#### Parsing
|
||
|
|
|
||
|
|
An implementation of the parser is shown below:
|
||
|
|
|
||
|
|
```c++
|
||
|
|
/// Parse an instance of a type registered to the toy dialect.
|
||
|
|
mlir::Type ToyDialect::parseType(mlir::DialectAsmParser &parser) const {
|
||
|
|
// Parse a struct type in the following form:
|
||
|
|
// struct-type ::= `struct` `<` type (`,` type)* `>`
|
||
|
|
|
||
|
|
// NOTE: All MLIR parser function return a ParseResult. This is a
|
||
|
|
// specialization of LogicalResult that auto-converts to a `true` boolean
|
||
|
|
// value on failure to allow for chaining, but may be used with explicit
|
||
|
|
// `mlir::failed/mlir::succeeded` as desired.
|
||
|
|
|
||
|
|
// Parse: `struct` `<`
|
||
|
|
if (parser.parseKeyword("struct") || parser.parseLess())
|
||
|
|
return Type();
|
||
|
|
|
||
|
|
// Parse the element types of the struct.
|
||
|
|
SmallVector<mlir::Type, 1> elementTypes;
|
||
|
|
do {
|
||
|
|
// Parse the current element type.
|
||
|
|
SMLoc typeLoc = parser.getCurrentLocation();
|
||
|
|
mlir::Type elementType;
|
||
|
|
if (parser.parseType(elementType))
|
||
|
|
return nullptr;
|
||
|
|
|
||
|
|
// Check that the type is either a TensorType or another StructType.
|
||
|
|
if (!elementType.isa<mlir::TensorType, StructType>()) {
|
||
|
|
parser.emitError(typeLoc, "element type for a struct must either "
|
||
|
|
"be a TensorType or a StructType, got: ")
|
||
|
|
<< elementType;
|
||
|
|
return Type();
|
||
|
|
}
|
||
|
|
elementTypes.push_back(elementType);
|
||
|
|
|
||
|
|
// Parse the optional: `,`
|
||
|
|
} while (succeeded(parser.parseOptionalComma()));
|
||
|
|
|
||
|
|
// Parse: `>`
|
||
|
|
if (parser.parseGreater())
|
||
|
|
return Type();
|
||
|
|
return StructType::get(elementTypes);
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
#### Printing
|
||
|
|
|
||
|
|
An implementation of the printer is shown below:
|
||
|
|
|
||
|
|
```c++
|
||
|
|
/// Print an instance of a type registered to the toy dialect.
|
||
|
|
void ToyDialect::printType(mlir::Type type,
|
||
|
|
mlir::DialectAsmPrinter &printer) const {
|
||
|
|
// Currently the only toy type is a struct type.
|
||
|
|
StructType structType = type.cast<StructType>();
|
||
|
|
|
||
|
|
// Print the struct type according to the parser format.
|
||
|
|
printer << "struct<";
|
||
|
|
llvm::interleaveComma(structType.getElementTypes(), printer);
|
||
|
|
printer << '>';
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
Before moving on, let's look at a quick of example showcasing the functionality
|
||
|
|
we have now:
|
||
|
|
|
||
|
|
```toy
|
||
|
|
struct Struct {
|
||
|
|
var a;
|
||
|
|
var b;
|
||
|
|
}
|
||
|
|
|
||
|
|
def multiply_transpose(Struct value) {
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
Which generates the following:
|
||
|
|
|
||
|
|
```mlir
|
||
|
|
module {
|
||
|
|
toy.func @multiply_transpose(%arg0: !toy.struct<tensor<*xf64>, tensor<*xf64>>) {
|
||
|
|
toy.return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
### Operating on `StructType`
|
||
|
|
|
||
|
|
Now that the `struct` type has been defined, and we can round-trip it through
|
||
|
|
the IR. The next step is to add support for using it within our operations.
|
||
|
|
|
||
|
|
#### Updating Existing Operations
|
||
|
|
|
||
|
|
A few of our existing operations, e.g. `ReturnOp`, will need to be updated to
|
||
|
|
handle `Toy_StructType`.
|
||
|
|
|
||
|
|
```tablegen
|
||
|
|
def ReturnOp : Toy_Op<"return", [Terminator, HasParent<"FuncOp">]> {
|
||
|
|
...
|
||
|
|
let arguments = (ins Variadic<Toy_Type>:$input);
|
||
|
|
...
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
#### Adding New `Toy` Operations
|
||
|
|
|
||
|
|
In addition to the existing operations, we will be adding a few new operations
|
||
|
|
that will provide more specific handling of `structs`.
|
||
|
|
|
||
|
|
##### `toy.struct_constant`
|
||
|
|
|
||
|
|
This new operation materializes a constant value for a struct. In our current
|
||
|
|
modeling, we just use an [array attribute](../../Dialects/Builtin.md/#arrayattr)
|
||
|
|
that contains a set of constant values for each of the `struct` elements.
|
||
|
|
|
||
|
|
```mlir
|
||
|
|
%0 = toy.struct_constant [
|
||
|
|
dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf64>
|
||
|
|
] : !toy.struct<tensor<*xf64>>
|
||
|
|
```
|
||
|
|
|
||
|
|
##### `toy.struct_access`
|
||
|
|
|
||
|
|
This new operation materializes the Nth element of a `struct` value.
|
||
|
|
|
||
|
|
```mlir
|
||
|
|
// Using %0 from above
|
||
|
|
%1 = toy.struct_access %0[0] : !toy.struct<tensor<*xf64>> -> tensor<*xf64>
|
||
|
|
```
|
||
|
|
|
||
|
|
With these operations, we can revisit our original example:
|
||
|
|
|
||
|
|
```toy
|
||
|
|
struct Struct {
|
||
|
|
var a;
|
||
|
|
var b;
|
||
|
|
}
|
||
|
|
|
||
|
|
# User defined generic function may operate on struct types as well.
|
||
|
|
def multiply_transpose(Struct value) {
|
||
|
|
# We can access the elements of a struct via the '.' operator.
|
||
|
|
return transpose(value.a) * transpose(value.b);
|
||
|
|
}
|
||
|
|
|
||
|
|
def main() {
|
||
|
|
# We initialize struct values using a composite initializer.
|
||
|
|
Struct value = {[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]};
|
||
|
|
|
||
|
|
# We pass these arguments to functions like we do with variables.
|
||
|
|
var c = multiply_transpose(value);
|
||
|
|
print(c);
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
and finally get a full MLIR module:
|
||
|
|
|
||
|
|
```mlir
|
||
|
|
module {
|
||
|
|
toy.func @multiply_transpose(%arg0: !toy.struct<tensor<*xf64>, tensor<*xf64>>) -> tensor<*xf64> {
|
||
|
|
%0 = toy.struct_access %arg0[0] : !toy.struct<tensor<*xf64>, tensor<*xf64>> -> tensor<*xf64>
|
||
|
|
%1 = toy.transpose(%0 : tensor<*xf64>) to tensor<*xf64>
|
||
|
|
%2 = toy.struct_access %arg0[1] : !toy.struct<tensor<*xf64>, tensor<*xf64>> -> tensor<*xf64>
|
||
|
|
%3 = toy.transpose(%2 : tensor<*xf64>) to tensor<*xf64>
|
||
|
|
%4 = toy.mul %1, %3 : tensor<*xf64>
|
||
|
|
toy.return %4 : tensor<*xf64>
|
||
|
|
}
|
||
|
|
toy.func @main() {
|
||
|
|
%0 = toy.struct_constant [
|
||
|
|
dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>,
|
||
|
|
dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
|
||
|
|
] : !toy.struct<tensor<*xf64>, tensor<*xf64>>
|
||
|
|
%1 = toy.generic_call @multiply_transpose(%0) : (!toy.struct<tensor<*xf64>, tensor<*xf64>>) -> tensor<*xf64>
|
||
|
|
toy.print %1 : tensor<*xf64>
|
||
|
|
toy.return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
#### Optimizing Operations on `StructType`
|
||
|
|
|
||
|
|
Now that we have a few operations operating on `StructType`, we also have many
|
||
|
|
new constant folding opportunities.
|
||
|
|
|
||
|
|
After inlining, the MLIR module in the previous section looks something like:
|
||
|
|
|
||
|
|
```mlir
|
||
|
|
module {
|
||
|
|
toy.func @main() {
|
||
|
|
%0 = toy.struct_constant [
|
||
|
|
dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>,
|
||
|
|
dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
|
||
|
|
] : !toy.struct<tensor<*xf64>, tensor<*xf64>>
|
||
|
|
%1 = toy.struct_access %0[0] : !toy.struct<tensor<*xf64>, tensor<*xf64>> -> tensor<*xf64>
|
||
|
|
%2 = toy.transpose(%1 : tensor<*xf64>) to tensor<*xf64>
|
||
|
|
%3 = toy.struct_access %0[1] : !toy.struct<tensor<*xf64>, tensor<*xf64>> -> tensor<*xf64>
|
||
|
|
%4 = toy.transpose(%3 : tensor<*xf64>) to tensor<*xf64>
|
||
|
|
%5 = toy.mul %2, %4 : tensor<*xf64>
|
||
|
|
toy.print %5 : tensor<*xf64>
|
||
|
|
toy.return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
We have several `toy.struct_access` operations that access into a
|
||
|
|
`toy.struct_constant`. As detailed in [chapter 3](Ch-3.md) (FoldConstantReshape),
|
||
|
|
we can add folders for these `toy` operations by setting the `hasFolder` bit
|
||
|
|
on the operation definition and providing a definition of the `*Op::fold`
|
||
|
|
method.
|
||
|
|
|
||
|
|
```c++
|
||
|
|
/// Fold constants.
|
||
|
|
OpFoldResult ConstantOp::fold(FoldAdaptor adaptor) { return value(); }
|
||
|
|
|
||
|
|
/// Fold struct constants.
|
||
|
|
OpFoldResult StructConstantOp::fold(FoldAdaptor adaptor) {
|
||
|
|
return value();
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Fold simple struct access operations that access into a constant.
|
||
|
|
OpFoldResult StructAccessOp::fold(FoldAdaptor adaptor) {
|
||
|
|
auto structAttr = adaptor.getInput().dyn_cast_or_null<mlir::ArrayAttr>();
|
||
|
|
if (!structAttr)
|
||
|
|
return nullptr;
|
||
|
|
|
||
|
|
size_t elementIndex = index().getZExtValue();
|
||
|
|
return structAttr[elementIndex];
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
To ensure that MLIR generates the proper constant operations when folding our
|
||
|
|
`Toy` operations, i.e. `ConstantOp` for `TensorType` and `StructConstant` for
|
||
|
|
`StructType`, we will need to provide an override for the dialect hook
|
||
|
|
`materializeConstant`. This allows for generic MLIR operations to create
|
||
|
|
constants for the `Toy` dialect when necessary.
|
||
|
|
|
||
|
|
```c++
|
||
|
|
mlir::Operation *ToyDialect::materializeConstant(mlir::OpBuilder &builder,
|
||
|
|
mlir::Attribute value,
|
||
|
|
mlir::Type type,
|
||
|
|
mlir::Location loc) {
|
||
|
|
if (type.isa<StructType>())
|
||
|
|
return builder.create<StructConstantOp>(loc, type,
|
||
|
|
value.cast<mlir::ArrayAttr>());
|
||
|
|
return builder.create<ConstantOp>(loc, type,
|
||
|
|
value.cast<mlir::DenseElementsAttr>());
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
With this, we can now generate code that can be generated to LLVM without any
|
||
|
|
changes to our pipeline.
|
||
|
|
|
||
|
|
```mlir
|
||
|
|
module {
|
||
|
|
toy.func @main() {
|
||
|
|
%0 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
|
||
|
|
%1 = toy.transpose(%0 : tensor<2x3xf64>) to tensor<3x2xf64>
|
||
|
|
%2 = toy.mul %1, %1 : tensor<3x2xf64>
|
||
|
|
toy.print %2 : tensor<3x2xf64>
|
||
|
|
toy.return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
You can build `toyc-ch7` and try yourself: `toyc-ch7
|
||
|
|
test/Examples/Toy/Ch7/struct-codegen.toy -emit=mlir`. More details on defining
|
||
|
|
custom types can be found in
|
||
|
|
[DefiningAttributesAndTypes](../../DefiningDialects/AttributesAndTypes.md).
|