feat: toy tutorial chapter 2.
This commit is contained in:
313
lib/Dialect.cpp
Normal file
313
lib/Dialect.cpp
Normal file
@@ -0,0 +1,313 @@
|
||||
//
|
||||
// Created by ricardo on 29/05/25.
|
||||
//
|
||||
|
||||
#include "Dialect.h"
|
||||
#include "hello/Dialect.cpp.inc"
|
||||
|
||||
#include <mlir/Interfaces/FunctionImplementation.h>
|
||||
|
||||
void mlir::hello::HelloDialect::initialize()
|
||||
{
|
||||
addOperations<
|
||||
#define GET_OP_LIST
|
||||
#include "hello/Ops.cpp.inc"
|
||||
>();
|
||||
}
|
||||
|
||||
using namespace mlir;
|
||||
using namespace mlir::hello;
|
||||
|
||||
|
||||
/// A generalized parser for binary operations. This parses the different forms
|
||||
/// of 'printBinaryOp' below.
|
||||
static ParseResult parseBinaryOp(OpAsmParser& parser,
|
||||
OperationState& result)
|
||||
{
|
||||
SmallVector<OpAsmParser::UnresolvedOperand, 2> operands;
|
||||
llvm::SMLoc operandsLoc = parser.getCurrentLocation();
|
||||
mlir::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 (mlir::FunctionType funcType = llvm::dyn_cast<mlir::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.
|
||||
mlir::Type resultType = *op->result_type_begin();
|
||||
if (llvm::all_of(op->getOperandTypes(),
|
||||
[=](mlir::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 mlir::hello::ConstantOp::build(OpBuilder& builder, OperationState& state,
|
||||
double value)
|
||||
{
|
||||
auto dataType = RankedTensorType::get({}, builder.getF64Type());
|
||||
auto dataAttribute = DenseElementsAttr::get(dataType, value);
|
||||
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();
|
||||
}
|
||||
|
||||
/// Verifier for the constant operation. This corresponds to the
|
||||
/// `let hasVerifier = 1` in the op definition.
|
||||
mlir::LogicalResult ConstantOp::verify()
|
||||
{
|
||||
// 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>(getResult().getType());
|
||||
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>(getValue().getType());
|
||||
if (attrType.getRank() != resultType.getRank())
|
||||
{
|
||||
return 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 emitOpError(
|
||||
"return type shape mismatches its attribute at dimension ")
|
||||
<< dim << ": " << attrType.getShape()[dim]
|
||||
<< " != " << resultType.getShape()[dim];
|
||||
}
|
||||
}
|
||||
return mlir::success();
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// 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); }
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// GenericCallOp
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
void GenericCallOp::build(mlir::OpBuilder& builder, mlir::OperationState& state,
|
||||
StringRef callee, ArrayRef<mlir::Value> arguments)
|
||||
{
|
||||
// Generic call always returns an unranked Tensor initially.
|
||||
state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
|
||||
state.addOperands(arguments);
|
||||
state.addAttribute("callee",
|
||||
mlir::SymbolRefAttr::get(builder.getContext(), callee));
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// 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(OpAsmParser& parser,
|
||||
OperationState& result)
|
||||
{
|
||||
// Dispatch to the FunctionOpInterface provided utility method that parses the
|
||||
// function operation.
|
||||
auto buildFuncType =
|
||||
[](Builder& builder, ArrayRef<Type> argTypes,
|
||||
ArrayRef<Type> results,
|
||||
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());
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// 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); }
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// ReturnOp
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
mlir::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
|
||||
<< ")";
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// TransposeOp
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
void TransposeOp::build(mlir::OpBuilder& builder, mlir::OperationState& state,
|
||||
mlir::Value value)
|
||||
{
|
||||
state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
|
||||
state.addOperands(value);
|
||||
}
|
||||
|
||||
mlir::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();
|
||||
}
|
||||
|
||||
|
||||
#define GET_OP_CLASSES
|
||||
#include "hello/Ops.cpp.inc"
|
||||
Reference in New Issue
Block a user