Files
TheRock/compiler/amd-llvm/mlir/examples/toy/Ch7/mlir/Dialect.cpp
T
derek.holloway 7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
Flatten subprojects into a single base monorepo
2026-09-02 16:34:47 -07:00

666 lines
27 KiB
C++

//===- Dialect.cpp - Toy IR Dialect registration in MLIR ------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements the dialect for the Toy IR: custom type parsing and
// operation verification.
//
//===----------------------------------------------------------------------===//
#include "toy/Dialect.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/TypeSupport.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Interfaces/CallInterfaces.h"
#include "mlir/Interfaces/FunctionImplementation.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/InliningUtils.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/Hashing.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <string>
using namespace mlir;
using namespace mlir::toy;
#include "toy/Dialect.cpp.inc"
//===----------------------------------------------------------------------===//
// ToyInlinerInterface
//===----------------------------------------------------------------------===//
/// This class defines the interface for handling inlining with Toy
/// operations.
struct ToyInlinerInterface : public DialectInlinerInterface {
using DialectInlinerInterface::DialectInlinerInterface;
//===--------------------------------------------------------------------===//
// Analysis Hooks
//===--------------------------------------------------------------------===//
/// All call operations within toy can be inlined.
bool isLegalToInline(Operation *call, Operation *callable,
bool wouldBeCloned) const final {
return true;
}
/// All operations within toy can be inlined.
bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
return true;
}
// All functions within toy can be inlined.
bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
return true;
}
//===--------------------------------------------------------------------===//
// Transformation Hooks
//===--------------------------------------------------------------------===//
/// Handle the given inlined terminator(toy.return) by replacing it with a new
/// operation as necessary.
void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {
// Only "toy.return" needs to be handled here.
auto returnOp = cast<ReturnOp>(op);
// Replace the values directly with the return operands.
assert(returnOp.getNumOperands() == valuesToRepl.size());
for (const auto &it : llvm::enumerate(returnOp.getOperands()))
valuesToRepl[it.index()].replaceAllUsesWith(it.value());
}
/// Attempts to materialize a conversion for a type mismatch between a call
/// from this dialect, and a callable region. This method should generate an
/// operation that takes 'input' as the only operand, and produces a single
/// result of 'resultType'. If a conversion can not be generated, nullptr
/// should be returned.
Operation *materializeCallConversion(OpBuilder &builder, Value input,
Type resultType,
Location conversionLoc) const final {
return CastOp::create(builder, conversionLoc, resultType, input);
}
};
//===----------------------------------------------------------------------===//
// Toy Operations
//===----------------------------------------------------------------------===//
/// A generalized parser for binary operations. This parses the different forms
/// of 'printBinaryOp' below.
static mlir::ParseResult parseBinaryOp(mlir::OpAsmParser &parser,
mlir::OperationState &result) {
SmallVector<mlir::OpAsmParser::UnresolvedOperand, 2> operands;
SMLoc operandsLoc = parser.getCurrentLocation();
Type type;
if (parser.parseOperandList(operands, /*requiredOperandCount=*/2) ||
parser.parseOptionalAttrDict(result.attributes) ||
parser.parseColonType(type))
return mlir::failure();
// If the type is a function type, it contains the input and result types of
// this operation.
if (FunctionType funcType = llvm::dyn_cast<FunctionType>(type)) {
if (parser.resolveOperands(operands, funcType.getInputs(), operandsLoc,
result.operands))
return mlir::failure();
result.addTypes(funcType.getResults());
return mlir::success();
}
// Otherwise, the parsed type is the type of both operands and results.
if (parser.resolveOperands(operands, type, result.operands))
return mlir::failure();
result.addTypes(type);
return mlir::success();
}
/// A generalized printer for binary operations. It prints in two different
/// forms depending on if all of the types match.
static void printBinaryOp(mlir::OpAsmPrinter &printer, mlir::Operation *op) {
printer << " " << op->getOperands();
printer.printOptionalAttrDict(op->getAttrs());
printer << " : ";
// If all of the types are the same, print the type directly.
Type resultType = *op->result_type_begin();
if (llvm::all_of(op->getOperandTypes(),
[=](Type type) { return type == resultType; })) {
printer << resultType;
return;
}
// Otherwise, print a functional type.
printer.printFunctionalType(op->getOperandTypes(), op->getResultTypes());
}
//===----------------------------------------------------------------------===//
// ConstantOp
//===----------------------------------------------------------------------===//
/// Build a constant operation.
/// The builder is passed as an argument, so is the state that this method is
/// expected to fill in order to build the operation.
void ConstantOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
double value) {
auto dataType = RankedTensorType::get({}, builder.getF64Type());
auto dataAttribute = DenseElementsAttr::get(dataType, value);
ConstantOp::build(builder, state, dataType, dataAttribute);
}
/// The 'OpAsmParser' class provides a collection of methods for parsing
/// various punctuation, as well as attributes, operands, types, etc. Each of
/// these methods returns a `ParseResult`. This class is a wrapper around
/// `LogicalResult` that can be converted to a boolean `true` value on failure,
/// or `false` on success. This allows for easily chaining together a set of
/// parser rules. These rules are used to populate an `mlir::OperationState`
/// similarly to the `build` methods described above.
mlir::ParseResult ConstantOp::parse(mlir::OpAsmParser &parser,
mlir::OperationState &result) {
mlir::DenseElementsAttr value;
if (parser.parseOptionalAttrDict(result.attributes) ||
parser.parseAttribute(value, "value", result.attributes))
return failure();
result.addTypes(value.getType());
return success();
}
/// The 'OpAsmPrinter' class is a stream that allows for formatting
/// strings, attributes, operands, types, etc.
void ConstantOp::print(mlir::OpAsmPrinter &printer) {
printer << " ";
printer.printOptionalAttrDict((*this)->getAttrs(), /*elidedAttrs=*/{"value"});
printer << getValue();
}
/// Verify that the given attribute value is valid for the given type.
static llvm::LogicalResult verifyConstantForType(mlir::Type type,
mlir::Attribute opaqueValue,
mlir::Operation *op) {
if (llvm::isa<mlir::TensorType>(type)) {
// Check that the value is an elements attribute.
auto attrValue = llvm::dyn_cast<mlir::DenseFPElementsAttr>(opaqueValue);
if (!attrValue)
return op->emitError("constant of TensorType must be initialized by "
"a DenseFPElementsAttr, got ")
<< opaqueValue;
// If the return type of the constant is not an unranked tensor, the shape
// must match the shape of the attribute holding the data.
auto resultType = llvm::dyn_cast<mlir::RankedTensorType>(type);
if (!resultType)
return success();
// Check that the rank of the attribute type matches the rank of the
// constant result type.
auto attrType = llvm::cast<mlir::RankedTensorType>(attrValue.getType());
if (attrType.getRank() != resultType.getRank()) {
return op->emitOpError("return type must match the one of the attached "
"value attribute: ")
<< attrType.getRank() << " != " << resultType.getRank();
}
// Check that each of the dimensions match between the two types.
for (int dim = 0, dimE = attrType.getRank(); dim < dimE; ++dim) {
if (attrType.getShape()[dim] != resultType.getShape()[dim]) {
return op->emitOpError(
"return type shape mismatches its attribute at dimension ")
<< dim << ": " << attrType.getShape()[dim]
<< " != " << resultType.getShape()[dim];
}
}
return mlir::success();
}
auto resultType = llvm::cast<StructType>(type);
llvm::ArrayRef<mlir::Type> resultElementTypes = resultType.getElementTypes();
// Verify that the initializer is an Array.
auto attrValue = llvm::dyn_cast<ArrayAttr>(opaqueValue);
if (!attrValue || attrValue.getValue().size() != resultElementTypes.size())
return op->emitError("constant of StructType must be initialized by an "
"ArrayAttr with the same number of elements, got ")
<< opaqueValue;
// Check that each of the elements are valid.
llvm::ArrayRef<mlir::Attribute> attrElementValues = attrValue.getValue();
for (const auto it : llvm::zip(resultElementTypes, attrElementValues))
if (failed(verifyConstantForType(std::get<0>(it), std::get<1>(it), op)))
return mlir::failure();
return mlir::success();
}
/// Verifier for the constant operation. This corresponds to the `::verify(...)`
/// in the op definition.
llvm::LogicalResult ConstantOp::verify() {
return verifyConstantForType(getResult().getType(), getValue(), *this);
}
llvm::LogicalResult StructConstantOp::verify() {
return verifyConstantForType(getResult().getType(), getValue(), *this);
}
/// Infer the output shape of the ConstantOp, this is required by the shape
/// inference interface.
void ConstantOp::inferShapes() {
getResult().setType(cast<TensorType>(getValue().getType()));
}
//===----------------------------------------------------------------------===//
// AddOp
//===----------------------------------------------------------------------===//
void AddOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
mlir::Value lhs, mlir::Value rhs) {
state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
state.addOperands({lhs, rhs});
}
mlir::ParseResult AddOp::parse(mlir::OpAsmParser &parser,
mlir::OperationState &result) {
return parseBinaryOp(parser, result);
}
void AddOp::print(mlir::OpAsmPrinter &p) { printBinaryOp(p, *this); }
/// Infer the output shape of the AddOp, this is required by the shape inference
/// interface.
void AddOp::inferShapes() { getResult().setType(getLhs().getType()); }
//===----------------------------------------------------------------------===//
// CastOp
//===----------------------------------------------------------------------===//
/// Infer the output shape of the CastOp, this is required by the shape
/// inference interface.
void CastOp::inferShapes() { getResult().setType(getInput().getType()); }
/// Returns true if the given set of input and result types are compatible with
/// this cast operation. This is required by the `CastOpInterface` to verify
/// this operation and provide other additional utilities.
bool CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
if (inputs.size() != 1 || outputs.size() != 1)
return false;
// The inputs must be Tensors with the same element type.
TensorType input = llvm::dyn_cast<TensorType>(inputs.front());
TensorType output = llvm::dyn_cast<TensorType>(outputs.front());
if (!input || !output || input.getElementType() != output.getElementType())
return false;
// The shape is required to match if both types are ranked.
return !input.hasRank() || !output.hasRank() || input == output;
}
//===----------------------------------------------------------------------===//
// FuncOp
//===----------------------------------------------------------------------===//
void FuncOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
llvm::StringRef name, mlir::FunctionType type,
llvm::ArrayRef<mlir::NamedAttribute> attrs) {
// FunctionOpInterface provides a convenient `build` method that will populate
// the state of our FuncOp, and create an entry block.
buildWithEntryBlock(builder, state, name, type, attrs, type.getInputs());
}
mlir::ParseResult FuncOp::parse(mlir::OpAsmParser &parser,
mlir::OperationState &result) {
// Dispatch to the FunctionOpInterface provided utility method that parses the
// function operation.
auto buildFuncType =
[](mlir::Builder &builder, llvm::ArrayRef<mlir::Type> argTypes,
llvm::ArrayRef<mlir::Type> results,
mlir::function_interface_impl::VariadicFlag,
std::string &) { return builder.getFunctionType(argTypes, results); };
return mlir::function_interface_impl::parseFunctionOp(
parser, result, /*allowVariadic=*/false,
getFunctionTypeAttrName(result.name), buildFuncType,
getArgAttrsAttrName(result.name), getResAttrsAttrName(result.name));
}
void FuncOp::print(mlir::OpAsmPrinter &p) {
// Dispatch to the FunctionOpInterface provided utility method that prints the
// function operation.
mlir::function_interface_impl::printFunctionOp(
p, *this, /*isVariadic=*/false, getFunctionTypeAttrName(),
getArgAttrsAttrName(), getResAttrsAttrName());
}
//===----------------------------------------------------------------------===//
// GenericCallOp
//===----------------------------------------------------------------------===//
void GenericCallOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
mlir::Type resultType, StringRef callee,
ArrayRef<mlir::Value> arguments) {
state.addTypes(resultType);
state.addOperands(arguments);
state.addAttribute("callee",
mlir::SymbolRefAttr::get(builder.getContext(), callee));
}
/// Return the callee of the generic call operation, this is required by the
/// call interface.
CallInterfaceCallable GenericCallOp::getCallableForCallee() {
return (*this)->getAttrOfType<SymbolRefAttr>("callee");
}
/// Set the callee for the generic call operation, this is required by the call
/// interface.
void GenericCallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
(*this)->setAttr("callee", cast<SymbolRefAttr>(callee));
}
/// Get the argument operands to the called function, this is required by the
/// call interface.
Operation::operand_range GenericCallOp::getArgOperands() { return getInputs(); }
/// Get the argument operands to the called function as a mutable range, this is
/// required by the call interface.
MutableOperandRange GenericCallOp::getArgOperandsMutable() {
return getInputsMutable();
}
//===----------------------------------------------------------------------===//
// MulOp
//===----------------------------------------------------------------------===//
void MulOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
mlir::Value lhs, mlir::Value rhs) {
state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
state.addOperands({lhs, rhs});
}
mlir::ParseResult MulOp::parse(mlir::OpAsmParser &parser,
mlir::OperationState &result) {
return parseBinaryOp(parser, result);
}
void MulOp::print(mlir::OpAsmPrinter &p) { printBinaryOp(p, *this); }
/// Infer the output shape of the MulOp, this is required by the shape inference
/// interface.
void MulOp::inferShapes() { getResult().setType(getLhs().getType()); }
//===----------------------------------------------------------------------===//
// ReturnOp
//===----------------------------------------------------------------------===//
llvm::LogicalResult ReturnOp::verify() {
// We know that the parent operation is a function, because of the 'HasParent'
// trait attached to the operation definition.
auto function = cast<FuncOp>((*this)->getParentOp());
/// ReturnOps can only have a single optional operand.
if (getNumOperands() > 1)
return emitOpError() << "expects at most 1 return operand";
// The operand number and types must match the function signature.
const auto &results = function.getFunctionType().getResults();
if (getNumOperands() != results.size())
return emitOpError() << "does not return the same number of values ("
<< getNumOperands() << ") as the enclosing function ("
<< results.size() << ")";
// If the operation does not have an input, we are done.
if (!hasOperand())
return mlir::success();
auto inputType = *operand_type_begin();
auto resultType = results.front();
// Check that the result type of the function matches the operand type.
if (inputType == resultType ||
llvm::isa<mlir::UnrankedTensorType>(inputType) ||
llvm::isa<mlir::UnrankedTensorType>(resultType))
return mlir::success();
return emitError() << "type of return operand (" << inputType
<< ") doesn't match function result type (" << resultType
<< ")";
}
//===----------------------------------------------------------------------===//
// StructAccessOp
//===----------------------------------------------------------------------===//
void StructAccessOp::build(mlir::OpBuilder &b, mlir::OperationState &state,
mlir::Value input, size_t index) {
// Extract the result type from the input type.
StructType structTy = llvm::cast<StructType>(input.getType());
assert(index < structTy.getNumElementTypes());
mlir::Type resultType = structTy.getElementTypes()[index];
// Call into the auto-generated build method.
build(b, state, resultType, input, b.getI64IntegerAttr(index));
}
llvm::LogicalResult StructAccessOp::verify() {
StructType structTy = llvm::cast<StructType>(getInput().getType());
size_t indexValue = getIndex();
if (indexValue >= structTy.getNumElementTypes())
return emitOpError()
<< "index should be within the range of the input struct type";
mlir::Type resultType = getResult().getType();
if (resultType != structTy.getElementTypes()[indexValue])
return emitOpError() << "must have the same result type as the struct "
"element referred to by the index";
return mlir::success();
}
//===----------------------------------------------------------------------===//
// TransposeOp
//===----------------------------------------------------------------------===//
void TransposeOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
mlir::Value value) {
state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
state.addOperands(value);
}
void TransposeOp::inferShapes() {
auto arrayTy = llvm::cast<RankedTensorType>(getOperand().getType());
SmallVector<int64_t, 2> dims(llvm::reverse(arrayTy.getShape()));
getResult().setType(RankedTensorType::get(dims, arrayTy.getElementType()));
}
llvm::LogicalResult TransposeOp::verify() {
auto inputType = llvm::dyn_cast<RankedTensorType>(getOperand().getType());
auto resultType = llvm::dyn_cast<RankedTensorType>(getType());
if (!inputType || !resultType)
return mlir::success();
auto inputShape = inputType.getShape();
if (!std::equal(inputShape.begin(), inputShape.end(),
resultType.getShape().rbegin())) {
return emitError()
<< "expected result shape to be a transpose of the input";
}
return mlir::success();
}
//===----------------------------------------------------------------------===//
// Toy Types
//===----------------------------------------------------------------------===//
namespace mlir {
namespace toy {
namespace detail {
/// 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, see the `StructType::get` method.
/// 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.
/// 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;
};
} // namespace detail
} // namespace toy
} // namespace mlir
/// Create an instance of a `StructType` with the given element types. There
/// *must* be at least one element type.
StructType 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 the context 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> StructType::getElementTypes() {
// 'getImpl' returns a pointer to the internal storage instance.
return getImpl()->elementTypes;
}
/// 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 (!llvm::isa<mlir::TensorType, StructType>(elementType)) {
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);
}
/// 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 = llvm::cast<StructType>(type);
// Print the struct type according to the parser format.
printer << "struct<";
llvm::interleaveComma(structType.getElementTypes(), printer);
printer << '>';
}
//===----------------------------------------------------------------------===//
// TableGen'd op method definitions
//===----------------------------------------------------------------------===//
#define GET_OP_CLASSES
#include "toy/Ops.cpp.inc"
//===----------------------------------------------------------------------===//
// ToyDialect
//===----------------------------------------------------------------------===//
/// Dialect initialization, the instance will be owned by the context. This is
/// the point of registration of types and operations for the dialect.
void ToyDialect::initialize() {
addOperations<
#define GET_OP_LIST
#include "toy/Ops.cpp.inc"
>();
addInterfaces<ToyInlinerInterface>();
addTypes<StructType>();
}
mlir::Operation *ToyDialect::materializeConstant(mlir::OpBuilder &builder,
mlir::Attribute value,
mlir::Type type,
mlir::Location loc) {
if (llvm::isa<StructType>(type))
return StructConstantOp::create(builder, loc, type,
llvm::cast<mlir::ArrayAttr>(value));
return ConstantOp::create(builder, loc, type,
llvm::cast<mlir::DenseElementsAttr>(value));
}