Program Listing for File DFGOps.cpp

↰ Return to documentation for file (lib/Dialect/DFG/IR/DFGOps.cpp)

#include "laksa-mlir/Dialect/DFG/IR/DFGOps.h"

#include "laksa-mlir/Dialect/DFG/DFGEnums.h"
#include "laksa-mlir/Dialect/DFG/IR/DFGBase.h"
#include "laksa-mlir/Dialect/DFG/IR/DFGTypes.h"
#include "laksa-mlir/Dialect/DFG/Interfaces/NodeInterface.h"
#include "laksa-mlir/IR/LaksaAttributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/FunctionImplementation.h"
#include "mlir/Transforms/InliningUtils.h"

#include <algorithm>
#include <cstdint>
#include <iterator>
#include <llvm/ADT/ArrayRef.h>
#include <llvm/ADT/StringRef.h>
#include <llvm/IR/PassManager.h>
#include <llvm/Support/Debug.h>
#include <llvm/Support/LogicalResult.h>
#include <llvm/Support/SMLoc.h>
#include <mlir/IR/Attributes.h>
#include <mlir/IR/BuiltinAttributes.h>
#include <mlir/IR/BuiltinOps.h>
#include <mlir/IR/BuiltinTypeInterfaces.h>
#include <mlir/IR/Diagnostics.h>
#include <mlir/IR/IRMapping.h>
#include <mlir/IR/OpImplementation.h>
#include <mlir/IR/OperationSupport.h>
#include <mlir/IR/SymbolTable.h>
#include <mlir/IR/Types.h>
#include <mlir/IR/Value.h>
#include <mlir/IR/ValueRange.h>
#include <mlir/Support/WalkResult.h>
#include <mlir/Transforms/DialectConversion.h>
#include <optional>
#include <ranges>

using namespace mlir;
using namespace mlir::dfg;

//===- Generated implementation -------------------------------------------===//

#define GET_OP_CLASSES
#include "laksa-mlir/Dialect/DFG/IR/DFGOps.cpp.inc"

//===----------------------------------------------------------------------===//

// For multiple variadic attributes
constexpr char kOperandSegmentSizesAttr[] = "operandSegmentSizes";

static ParseResult parseChannelArgumentList(
    OpAsmParser &parser,
    SmallVectorImpl<OpAsmParser::Argument> &arguments)
{
    return parser.parseCommaSeparatedList(
        OpAsmParser::Delimiter::Paren,
        [&]() -> ParseResult {
            OpAsmParser::Argument argument;
            if (parser.parseArgument(
                    argument,
                    /*allowType=*/true,
                    /*allowAttrs=*/false))
                return failure();

            arguments.push_back(argument);
            return success();
        });
}

//===----------------------------------------------------------------------===//
// DFGDialect Node Operations
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// ProcessOp
//===----------------------------------------------------------------------===//

void ProcessOp::getAsmBlockArgumentNames(
    Region &region,
    OpAsmSetValueNameFn setNameFn)
{
    unsigned numInputs = getFunctionType().getNumInputs();
    for (unsigned i = 0, e = region.getNumArguments(); i < e; ++i)
        if (i < numInputs)
            setNameFn(region.getArgument(i), ("in" + std::to_string(i)));
        else
            setNameFn(
                region.getArgument(i),
                ("out" + std::to_string(i - numInputs)));
}

void ProcessOp::build(
    OpBuilder &builder,
    OperationState &state,
    StringRef name,
    FunctionType functionType,
    ArrayRef<int64_t> multiplicity,
    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder)
{
    state.addAttribute(
        SymbolTable::getSymbolAttrName(),
        builder.getStringAttr(name));

    // Get the function type with dfg types from base
    SmallVector<Type> inPortTypes, outPortTypes;
    for (auto inTy : functionType.getInputs())
        inPortTypes.push_back(
            isa<OutputType>(inTy)   ? inTy
            : isa<ShapedType>(inTy) ? OutputType::get(cast<ShapedType>(inTy))
                                    : OutputType::get(inTy));
    for (auto outTy : functionType.getResults())
        outPortTypes.push_back(
            isa<InputType>(outTy)    ? outTy
            : isa<ShapedType>(outTy) ? InputType::get(cast<ShapedType>(outTy))
                                     : InputType::get(outTy));
    auto dfgFunctionType =
        FunctionType::get(builder.getContext(), inPortTypes, outPortTypes);
    state.addAttribute(
        ProcessOp::getFunctionTypeAttrName(state.name),
        TypeAttr::get(dfgFunctionType));

    if (!multiplicity.empty())
        state.addAttribute(
            "multiplicity",
            DenseI64ArrayAttr::get(state.getContext(), multiplicity));

    OpBuilder::InsertionGuard guard(builder);
    Region* region = state.addRegion();
    auto regionBlock = builder.createBlock(region);

    // Create block arguments from the new function type
    SmallVector<Type> blockArgTypes;
    blockArgTypes.append(
        dfgFunctionType.getInputs().begin(),
        dfgFunctionType.getInputs().end());
    blockArgTypes.append(
        dfgFunctionType.getResults().begin(),
        dfgFunctionType.getResults().end());
    regionBlock->addArguments(
        blockArgTypes,
        SmallVector<Location>(blockArgTypes.size(), builder.getUnknownLoc()));

    if (bodyBuilder)
        bodyBuilder(builder, state.location, regionBlock->getArguments());
}

ParseResult ProcessOp::parse(OpAsmParser &parser, OperationState &state)
{
    auto &builder = parser.getBuilder();

    // Parse the process' name
    StringAttr nameAttr;
    if (parser.parseSymbolName(
            nameAttr,
            getSymNameAttrName(state.name),
            state.attributes))
        return failure();

    // Parse the signature of the process
    SmallVector<OpAsmParser::Argument> inputValues, outputValues;
    SMLoc signatureLocation = parser.getCurrentLocation();

    // Parse inputs/outputs separately for later distinction
    if (succeeded(parser.parseOptionalKeyword("inputs"))) {
        if (parseChannelArgumentList(parser, inputValues)) return failure();
    }

    if (succeeded(parser.parseOptionalKeyword("outputs"))) {
        if (parseChannelArgumentList(parser, outputValues)) return failure();
    }

    SmallVector<Type> inputTypes, outputTypes;
    inputTypes.reserve(inputValues.size());
    outputTypes.reserve(outputValues.size());

    for (auto &input : inputValues) inputTypes.push_back(input.type);
    for (auto &output : outputValues) outputTypes.push_back(output.type);
    Type funcType = builder.getFunctionType(inputTypes, outputTypes);

    if (!funcType) {
        return parser.emitError(signatureLocation)
               << "Failed to construct process type";
    }

    state.addAttribute(
        getFunctionTypeAttrName(state.name),
        TypeAttr::get(funcType));

    // Merge both argument lists for the block arguments
    inputValues.append(outputValues);

    OptionalParseResult attrResult =
        parser.parseOptionalAttrDictWithKeyword(state.attributes);
    if (attrResult.has_value() && failed(*attrResult)) return failure();

    // Parse the attached region, if any
    auto* body = state.addRegion();
    SMLoc loc = parser.getCurrentLocation();
    OptionalParseResult parseResult = parser.parseOptionalRegion(
        *body,
        inputValues,
        /*enableNameShadowing=*/false);

    if (parseResult.has_value()) {
        if (failed(*parseResult)) return failure();
        if (body->empty())
            return parser.emitError(loc, "expected non-empty process body");
    }

    return success();
}

void ProcessOp::print(OpAsmPrinter &p)
{
    Operation* op = getOperation();
    Region &body = getBody();
    bool isExternProcess = isExternal();

    // Print the operation and function name
    auto name = op->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName())
                    .getValue();

    p << ' ';
    p.printSymbolName(name);

    auto funcType = getFunctionType();
    auto inputTypes = funcType.getInputs();
    if (!inputTypes.empty()) {
        p << " inputs(";
        for (unsigned i = 0; i < inputTypes.size(); ++i) {
            if (i > 0) p << ", ";

            if (isExternProcess)
                p << "\%arg" << i;
            else
                p.printOperand(body.getArgument(i));
            p << ": " << inputTypes[i];
        }
        p << ")";
    }

    auto outputTypes = funcType.getResults();
    if (!outputTypes.empty()) {
        p << " outputs(";
        unsigned numInputs = inputTypes.size();
        for (unsigned i = 0; i < outputTypes.size(); ++i) {
            if (i > 0) p << ", ";

            if (isExternProcess)
                p << "\%arg" << i;
            else
                p.printOperand(body.getArgument(i + numInputs));
            p << ": " << outputTypes[i];
        }
        p << ")";
    }

    SmallVector<StringRef, 3> elidedAttrs = {
        getFunctionTypeAttrName(),
        getSymNameAttrName()};
    if (auto multiplicityAttr =
            cast<DenseI64ArrayAttr>(op->getAttr("multiplicity"))) {
        if (multiplicityAttr.empty()) elidedAttrs.push_back("multiplicity");
    }

    // Print any attributes in the attribute list into the dict
    if (!op->getAttrs().empty())
        p.printOptionalAttrDictWithKeyword(
            op->getAttrs(),
            /*elidedAttrs=*/elidedAttrs);

    // Print the region
    if (!isExternProcess) {
        p << ' ';
        p.printRegion(
            body,
            /*printEntryBlockArgs =*/false,
            /*printBlockTerminators =*/true);
    }
}

LogicalResult ProcessOp::verify()
{
    FunctionType funcType = getFunctionType();

    // All input ports must be dfg::OutputType and output ports dfg::InputType.
    if (!llvm::all_of(funcType.getInputs(), [](Type t) {
            return isa<OutputType>(t);
        }))
        return emitError("Input ports must be of type dfg::OutputType");
    if (!llvm::all_of(funcType.getResults(), [](Type t) {
            return isa<InputType>(t);
        }))
        return emitError("Output ports must be of type dfg::InputType");

    // If a multiplicity is defined, it must cover all arguments!
    ArrayRef<int64_t> multiplicity = getMultiplicity();
    size_t numArguments = funcType.getNumInputs() + funcType.getNumResults();
    if (!multiplicity.empty() && multiplicity.size() != numArguments)
        return emitError(
            "ProcessOp multiplicity must have a multiplicity for each "
            "channel");

    return success();
}

namespace {
struct EliminateDeadProcess final : public OpRewritePattern<ProcessOp> {
    using OpRewritePattern<ProcessOp>::OpRewritePattern;
    LogicalResult
    matchAndRewrite(ProcessOp op, PatternRewriter &rewriter) const override
    {
        auto module = op->getParentOfType<ModuleOp>();
        if (!module) return failure();

        auto processName = op.getNodeName();
        bool isInstantiated = false;
        module->walk([&](InstantiateOp instantiateOp) {
            if (instantiateOp.getNodeName() == processName)
                isInstantiated = true;
        });

        if (isInstantiated) return failure();

        rewriter.eraseOp(op);
        return success();
    }
};
} // namespace

void ProcessOp::getCanonicalizationPatterns(
    RewritePatternSet &patterns,
    MLIRContext* context)
{ patterns.add<EliminateDeadProcess>(context); }

//===----------------------------------------------------------------------===//
// LoopOp
//===----------------------------------------------------------------------===//

void LoopOp::build(
    OpBuilder &builder,
    OperationState &state,
    ValueRange inputChannels,
    ValueRange outputChannels,
    function_ref<void(OpBuilder &, Location)> bodyBuilder)
{
    assert(bodyBuilder && "loop must have contents");

    state.addOperands(inputChannels);
    state.addOperands(outputChannels);
    llvm::copy(
        ArrayRef<int32_t>(
            {static_cast<int32_t>(inputChannels.size()),
             static_cast<int32_t>(outputChannels.size())}),
        state.getOrAddProperties<Properties>().operandSegmentSizes.begin());

    // Add init body (no arguments)
    OpBuilder::InsertionGuard guard(builder);
    Region* region = state.addRegion();
    builder.createBlock(region);

    bodyBuilder(builder, state.location);
}

ParseResult LoopOp::parse(OpAsmParser &parser, OperationState &state)
{
    SmallVector<OpAsmParser::Argument> inputValues, outputValues;

    // Parse inputs
    SMLoc inputLocation = parser.getCurrentLocation();
    if (succeeded(parser.parseOptionalKeyword("inputs"))) {
        if (parseChannelArgumentList(parser, inputValues)) return failure();
    }
    int32_t numInputs = inputValues.size();
    // Parse outputs
    SMLoc outputLocation = parser.getCurrentLocation();
    if (succeeded(parser.parseOptionalKeyword("outputs"))) {
        if (parseChannelArgumentList(parser, outputValues)) return failure();
    }
    int32_t numOutputs = outputValues.size();

    // Operand and types
    SmallVector<OpAsmParser::UnresolvedOperand, 4> inputs, outputs;
    inputs.reserve(numInputs);
    outputs.reserve(numOutputs);
    SmallVector<Type> inputTypes, outputTypes;
    inputTypes.reserve(numInputs);
    outputTypes.reserve(numOutputs);

    for (auto &input : inputValues) {
        inputs.push_back(input.ssaName);
        inputTypes.push_back(input.type);
    }
    for (auto &output : outputValues) {
        outputs.push_back(output.ssaName);
        outputTypes.push_back(output.type);
    }

    // Resolve the inputs and outputs
    if (parser
            .resolveOperands(inputs, inputTypes, inputLocation, state.operands))
        return failure();
    if (parser.resolveOperands(
            outputs,
            outputTypes,
            outputLocation,
            state.operands))
        return failure();

    // Add derived `operand_segment_sizes` attribute based on parsed
    // operands.
    auto operandSegmentSizes =
        parser.getBuilder().getDenseI32ArrayAttr({numInputs, numOutputs});
    state.addAttribute(kOperandSegmentSizesAttr, operandSegmentSizes);

    // Parse the body region
    auto* body = state.addRegion();
    if (parser.parseRegion(*body, /*arguments=*/{})) return failure();

    return success();
}

void LoopOp::print(OpAsmPrinter &p)
{
    assert(
        !getOperands().empty()
        && "There has to be ports to be closed in LoopOp.");

    // Print inputs if exist
    auto inputs = getInputChannels();
    if (!inputs.empty()) {
        p << " inputs(";
        for (unsigned i = 0; i < inputs.size(); ++i) {
            if (i > 0) p << ", ";
            p.printOperand(inputs[i]);
            p << ": ";
            p.printType(inputs[i].getImpl()->getType());
        }
        p << ")";
    }

    // Print outputs if existent
    Operation::operand_range outputs = getOutputChannels();
    if (!outputs.empty()) {
        p << " outputs(";
        for (unsigned i = 0; i < outputs.size(); ++i) {
            if (i > 0) p << ", ";
            p.printOperand(outputs[i]);
            p << ": ";
            p.printType(outputs[i].getImpl()->getType());
        }
        p << ")";
    }

    // Print body region
    Region &body = getBody();
    if (!body.empty()) {
        p << ' ';
        p.printRegion(
            body,
            /*printEntryBlockArgs =*/false,
            /*printBlockTerminators =*/true);
    }
}

LogicalResult LoopOp::verify()
{
    auto loopFuncType = getFunctionType();
    auto loopNumInputs = loopFuncType.getNumInputs();
    auto loopNumOutputs = loopFuncType.getNumResults();

    auto processFuncType = getParentOp().getFunctionType();
    auto processNumInputs = processFuncType.getNumInputs();
    auto processNumOutputs = processFuncType.getNumResults();

    if (loopNumInputs != 0 && loopNumInputs != processNumInputs)
        return emitOpError("LoopOp should monitor all input ports or none.");

    if (loopNumOutputs != 0 && loopNumOutputs != processNumOutputs)
        return emitOpError("LoopOp should monitor all output ports or none.");

    return success();
}

//===----------------------------------------------------------------------===//
// OperatorOp
//===----------------------------------------------------------------------===//

void OperatorOp::getAsmBlockArgumentNames(
    Region &region,
    OpAsmSetValueNameFn setNameFn)
{
    unsigned numInputs = getFunctionType().getNumInputs();
    for (unsigned i = 0, e = region.getNumArguments(); i < e; ++i)
        if (i < numInputs)
            setNameFn(region.getArgument(i), ("in" + std::to_string(i)));
        else
            setNameFn(
                region.getArgument(i),
                ("out" + std::to_string(i - numInputs)));
}

void OperatorOp::build(
    OpBuilder &builder,
    OperationState &state,
    StringRef name,
    FunctionType functionType,
    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder)
{
    state.addAttribute(
        SymbolTable::getSymbolAttrName(),
        builder.getStringAttr(name));
    state.addAttribute(
        OperatorOp::getFunctionTypeAttrName(state.name),
        TypeAttr::get(functionType));

    // Add init body (no arguments)
    OpBuilder::InsertionGuard guard(builder);
    Region* region = state.addRegion();
    auto regionBlock = builder.createBlock(region);

    // Create block arguments from the new function type
    SmallVector<Type> blockArgTypes;
    blockArgTypes.append(
        functionType.getInputs().begin(),
        functionType.getInputs().end());
    regionBlock->addArguments(
        blockArgTypes,
        SmallVector<Location>(blockArgTypes.size(), builder.getUnknownLoc()));

    if (bodyBuilder)
        bodyBuilder(builder, state.location, regionBlock->getArguments());
}

ParseResult OperatorOp::parse(OpAsmParser &parser, OperationState &state)
{
    auto &builder = parser.getBuilder();

    // Parse the operator's name
    StringAttr nameAttr;
    if (parser.parseSymbolName(
            nameAttr,
            getSymNameAttrName(state.name),
            state.attributes))
        return failure();

    // Parse the signature of the operator
    SmallVector<OpAsmParser::Argument> inputValues, outputValues;
    SMLoc signatureLocation = parser.getCurrentLocation();

    // Parse inputs/outputs separately for later distinction
    if (succeeded(parser.parseOptionalKeyword("inputs"))) {
        if (parseChannelArgumentList(parser, inputValues)) return failure();
    }

    if (succeeded(parser.parseOptionalKeyword("outputs"))) {
        if (parseChannelArgumentList(parser, outputValues)) return failure();
    }

    SmallVector<Type> argTypes, resultTypes;
    argTypes.reserve(inputValues.size());
    resultTypes.reserve(outputValues.size());

    for (auto &input : inputValues) argTypes.push_back(input.type);
    for (auto &output : outputValues) resultTypes.push_back(output.type);
    Type funcType = builder.getFunctionType(argTypes, resultTypes);

    if (!funcType) {
        return parser.emitError(signatureLocation)
               << "Failed to construct process type";
    }

    state.addAttribute(
        getFunctionTypeAttrName(state.name),
        TypeAttr::get(funcType));

    OptionalParseResult attrResult =
        parser.parseOptionalAttrDictWithKeyword(state.attributes);
    if (attrResult.has_value() && failed(*attrResult)) return failure();

    // Parse the attached region, if any
    auto* body = state.addRegion();
    SMLoc loc = parser.getCurrentLocation();
    OptionalParseResult parseResult = parser.parseOptionalRegion(
        *body,
        inputValues,
        /*enableNameShadowing=*/false);

    if (parseResult.has_value()) {
        if (failed(*parseResult)) return failure();
        if (body->empty())
            return parser.emitError(loc, "expected non-empty operator body");
    }

    return success();
}

void OperatorOp::print(OpAsmPrinter &p)
{
    Operation* op = getOperation();
    Region &body = getBody();
    bool isExternOperator = isExternal();

    // Print the operation and its name
    auto name = op->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName())
                    .getValue();

    p << ' ';
    p.printSymbolName(name);

    // Print input list
    auto inputTypes = getInputPortTypes();
    if (!inputTypes.empty()) {
        p << " inputs(";
        for (unsigned i = 0; i < inputTypes.size(); ++i) {
            if (i > 0) p << ", ";

            if (isExternOperator)
                p << "\%arg" << i;
            else
                p.printOperand(body.getArgument(i));
            p << ": " << inputTypes[i];
        }
        p << ")";
    }

    // Print output list
    auto outputTypes = getOutputPortTypes();
    if (!outputTypes.empty()) {
        p << " outputs(";
        unsigned numInputs = inputTypes.size();
        for (unsigned i = 0; i < outputTypes.size(); ++i) {
            if (i > 0) p << ", ";

            p << "\%arg" << i + numInputs;

            p << ": " << outputTypes[i];
        }
        p << ")";
    }

    SmallVector<StringRef, 3> elidedAttrs = {
        getFunctionTypeAttrName(),
        getSymNameAttrName()};

    // Print any attributes in the attribute list into the dict
    if (!op->getAttrs().empty())
        p.printOptionalAttrDictWithKeyword(
            op->getAttrs(),
            /*elidedAttrs=*/elidedAttrs);

    // Print the region
    if (!isExternOperator) {
        p << ' ';
        p.printRegion(
            body,
            /*printEntryBlockArgs =*/false,
            /*printBlockTerminators =*/true);
    }
}

LogicalResult OperatorOp::verify()
{
    for (auto &opi : getBody().getOps()) {
        // Only allow OutputOp if the op is from this dialect
        if (opi.getDialect()->getNamespace() == "dfg") {
            if (!isa<OutputOp>(opi))
                return emitOpError(
                    "Only OutputOp from DFG dialect is allowed inside "
                    "OperatorOp.");
        }
    }
    return success();
}

namespace {
struct EliminateDeadOperator final : public OpRewritePattern<OperatorOp> {
    using OpRewritePattern<OperatorOp>::OpRewritePattern;
    LogicalResult
    matchAndRewrite(OperatorOp op, PatternRewriter &rewriter) const override
    {
        auto module = op->getParentOfType<ModuleOp>();
        if (!module) return failure();

        auto operatorName = op.getNodeName();
        bool isInstantiated = false;
        module->walk([&](InstantiateOp instantiateOp) {
            if (instantiateOp.getNodeName() == operatorName)
                isInstantiated = true;
        });

        if (isInstantiated) return failure();

        rewriter.eraseOp(op);
        return success();
    }
};
} // namespace

void OperatorOp::getCanonicalizationPatterns(
    RewritePatternSet &patterns,
    MLIRContext* context)
{ patterns.add<EliminateDeadOperator>(context); }

//===----------------------------------------------------------------------===//
// RegionOp
//===----------------------------------------------------------------------===//

void RegionOp::getAsmBlockArgumentNames(
    Region &region,
    OpAsmSetValueNameFn setNameFn)
{
    unsigned numInputs = getFunctionType().getNumInputs();
    for (unsigned i = 0, e = region.getNumArguments(); i < e; ++i)
        if (i < numInputs)
            setNameFn(region.getArgument(i), ("in" + std::to_string(i)));
        else
            setNameFn(
                region.getArgument(i),
                ("out" + std::to_string(i - numInputs)));
}

void RegionOp::build(
    OpBuilder &builder,
    OperationState &state,
    StringRef name,
    FunctionType functionType,
    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder)
{
    state.addAttribute(
        SymbolTable::getSymbolAttrName(),
        builder.getStringAttr(name));

    // Get the function type with dfg types from base
    SmallVector<Type> inPortTypes, outPortTypes;
    for (auto inTy : functionType.getInputs())
        inPortTypes.push_back(
            isa<OutputType>(inTy)   ? inTy
            : isa<ShapedType>(inTy) ? OutputType::get(cast<ShapedType>(inTy))
                                    : OutputType::get(inTy));
    for (auto outTy : functionType.getResults())
        outPortTypes.push_back(
            isa<InputType>(outTy)    ? outTy
            : isa<ShapedType>(outTy) ? InputType::get(cast<ShapedType>(outTy))
                                     : InputType::get(outTy));
    auto dfgFunctionType =
        FunctionType::get(builder.getContext(), inPortTypes, outPortTypes);
    state.addAttribute(
        RegionOp::getFunctionTypeAttrName(state.name),
        TypeAttr::get(dfgFunctionType));

    OpBuilder::InsertionGuard guard(builder);
    Region* region = state.addRegion();
    auto regionBlock = builder.createBlock(region);

    // Create block arguments from the new function type
    SmallVector<Type> blockArgTypes;
    blockArgTypes.append(
        dfgFunctionType.getInputs().begin(),
        dfgFunctionType.getInputs().end());
    blockArgTypes.append(
        dfgFunctionType.getResults().begin(),
        dfgFunctionType.getResults().end());
    regionBlock->addArguments(
        blockArgTypes,
        SmallVector<Location>(blockArgTypes.size(), builder.getUnknownLoc()));

    if (bodyBuilder)
        bodyBuilder(builder, state.location, regionBlock->getArguments());
}

ParseResult RegionOp::parse(OpAsmParser &parser, OperationState &state)
{
    auto &builder = parser.getBuilder();

    // Parse region name
    StringAttr nameAttr;
    if (parser.parseSymbolName(
            nameAttr,
            getSymNameAttrName(state.name),
            state.attributes))
        return failure();

    // Parse the signature of this region
    SmallVector<OpAsmParser::Argument> inputValues, outputValues;
    SMLoc signatureLocation = parser.getCurrentLocation();

    // Parse inputs/outputs separately for later distinction
    if (succeeded(parser.parseOptionalKeyword("inputs"))) {
        if (parseChannelArgumentList(parser, inputValues)) return failure();
    }
    if (succeeded(parser.parseOptionalKeyword("outputs"))) {
        if (parseChannelArgumentList(parser, outputValues)) return failure();
    }

    // Create the signature
    SmallVector<Type> inputTypes, outputTypes;
    inputTypes.reserve(inputValues.size());
    outputTypes.reserve(outputValues.size());

    for (auto &input : inputValues) inputTypes.push_back(input.type);
    for (auto &output : outputValues) outputTypes.push_back(output.type);
    Type funcType = builder.getFunctionType(inputTypes, outputTypes);
    if (!funcType) {
        return parser.emitError(signatureLocation)
               << "Failed to construct operator type";
    }
    state.addAttribute(
        getFunctionTypeAttrName(state.name),
        TypeAttr::get(funcType));

    // Merge both argument lists for the block arguments
    inputValues.append(outputValues);

    OptionalParseResult attrResult =
        parser.parseOptionalAttrDictWithKeyword(state.attributes);
    if (attrResult.has_value() && failed(*attrResult)) return failure();

    // Parse the attached region
    auto* body = state.addRegion();
    SMLoc loc = parser.getCurrentLocation();
    OptionalParseResult parseResult = parser.parseRegion(
        *body,
        inputValues,
        /*enableNameShadowing=*/false);

    if (parseResult.has_value()) {
        if (failed(*parseResult)) return failure();
        if (body->empty())
            return parser.emitError(loc, "expected non-empty operator body");
    }

    return success();
}

void RegionOp::print(OpAsmPrinter &p)
{
    Operation* op = getOperation();
    Region &body = op->getRegion(0);

    // Print the region's name
    auto name = op->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName())
                    .getValue();
    p << ' ';
    p.printSymbolName(name);

    // Print input list
    auto funcType = getFunctionType();
    auto inputTypes = funcType.getInputs();
    if (!inputTypes.empty()) {
        p << " inputs(";
        for (unsigned i = 0; i < inputTypes.size(); ++i) {
            if (i > 0) p << ", ";
            p.printOperand(body.getArgument(i));
            p << " : " << inputTypes[i];
        }
        p << ") ";
    }

    // Print output list
    auto outputTypes = funcType.getResults();
    if (!outputTypes.empty()) {
        p << " outputs(";
        unsigned numInputs = inputTypes.size();
        for (unsigned i = 0; i < outputTypes.size(); ++i) {
            if (i > 0) p << ", ";
            p.printOperand(body.getArgument(i + numInputs));
            p << " : " << outputTypes[i];
        }
        p << ") ";
    }

    // Print any attributes in the attribute list into the dict
    if (!op->getAttrs().empty())
        p.printOptionalAttrDictWithKeyword(
            op->getAttrs(),
            /*elidedAttrs=*/{getFunctionTypeAttrName(), getSymNameAttrName()});

    // Print the region
    p << ' ';
    p.printRegion(
        body,
        /*printEntryBlockArgs =*/false,
        /*printBlockTerminators =*/true);
}

LogicalResult RegionOp::verify()
{
    auto funcType = getFunctionType();
    // All input ports must be dfg::OutputType and output ports dfg::InputType.
    if (!llvm::all_of(funcType.getInputs(), [](Type t) {
            return isa<OutputType>(t);
        }))
        return emitError("Input ports must be of type dfg::OutputType");
    if (!llvm::all_of(funcType.getResults(), [](Type t) {
            return isa<InputType>(t);
        }))
        return emitError("Output ports must be of type dfg::InputType");
    // Check if the content operations are supported
    for (auto &opi : getBody().getOps())
        if (!isa<ChannelOp, InstantiateOp, EmbedOp>(opi))
            ::emitWarning(
                opi.getLoc(),
                "Unsupported op used in this region with name ")
                << opi.getName();

    // Check if all ports are connected
    for (auto [idx, arg] : llvm::enumerate(getBody().getArguments())) {
        if (arg.use_empty()) {
            if (idx <= getNumInputPorts()) {
                ::emitWarning(getLoc(), "Detecting dangling input port #")
                    << idx;
            } else {
                ::emitWarning(getLoc(), "Detecting dangling output port #")
                    << idx - getNumInputPorts();
            }
        } else if (!arg.hasOneUse()) {
            if (idx <= getNumInputPorts()) {
                ::emitWarning(getLoc(), "Detecting input port #")
                    << idx << " used more than once";
            } else {
                ::emitWarning(getLoc(), "Detecting output port #")
                    << idx - getNumInputPorts() << " used more than once";
            }
        }
    }

    return success();
}

namespace {
struct ReorderContentOperations final : public OpRewritePattern<RegionOp> {
    using OpRewritePattern<RegionOp>::OpRewritePattern;
    LogicalResult
    matchAndRewrite(RegionOp op, PatternRewriter &rewriter) const override
    {
        if (isReordered[op]) return success();

        auto loc = op.getLoc();
        auto funcTy = op.getFunctionType();

        // Create a new region with same name and signature
        auto newRegion = RegionOp::create(
            rewriter,
            loc,
            op.getNodeName(),
            funcTy,
            [&](OpBuilder &, Location, ValueRange regionBlkArgs) {
                IRMapping mapper;
                for (auto [oldArg, newArg] :
                     llvm::zip(op.getBody().getArguments(), regionBlkArgs)) {
                    mapper.map(oldArg, newArg);
                }

                // Save operations
                SmallVector<Operation*> channels, instances;
                for (auto &opOldRegion : op.getBody().getOps())
                    if (isa<ChannelOp>(opOldRegion))
                        channels.push_back(&opOldRegion);
                    else if (isa<InstantiateOp, EmbedOp>(opOldRegion))
                        instances.push_back(&opOldRegion);

                // Copy operations in order
                for (auto &channel : channels) rewriter.clone(*channel, mapper);
                for (auto &instance : instances)
                    rewriter.clone(*instance, mapper);
            });

        if (Attribute rootAttr = op->getAttr(laksa::kRootAttrName))
            newRegion->setAttr(laksa::kRootAttrName, rootAttr);

        // Save this new region as reordered
        rewriter.replaceOp(op, newRegion);
        isReordered[newRegion] = true;
        return success();
    }

private:
    mutable DenseMap<RegionOp, bool> isReordered;
};
} // namespace

void RegionOp::getCanonicalizationPatterns(
    RewritePatternSet &patterns,
    MLIRContext* context)
{ patterns.add<ReorderContentOperations>(context); }

//===----------------------------------------------------------------------===//
// DFGDialect Edge Operations
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// ChannelOp
//===----------------------------------------------------------------------===//

void ChannelOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn)
{
    Block* block = getOperation()->getBlock();
    int count = 0;

    for (auto &opi : *block) {
        if (&opi == getOperation()) break;
        if (opi.getName() == getOperation()->getName()) ++count;
    }

    setNameFn(getResult(0), "in_port_" + std::to_string(count));
    setNameFn(getResult(1), "out_port_" + std::to_string(count));
}

void ChannelOp::build(
    OpBuilder &builder,
    OperationState &state,
    ArrayRef<int64_t> shape,
    Type tokenType,
    std::optional<int32_t> bufferSize)
{
    state.addAttribute(
        getTokenTypeAttrName(state.name),
        TypeAttr::get(tokenType));
    if (bufferSize.has_value())
        state.addAttribute(
            getBufferSizeAttrName(state.name),
            builder.getI32IntegerAttr(bufferSize.value()));
    state.addTypes(InputType::get(shape, tokenType));
    state.addTypes(OutputType::get(shape, tokenType));
}

void ChannelOp::build(
    OpBuilder &builder,
    OperationState &state,
    Type type,
    std::optional<int32_t> bufferSize)
{
    if (auto shapedType = dyn_cast<ShapedType>(type)) {
        build(
            builder,
            state,
            shapedType.getShape(),
            shapedType.getElementType(),
            bufferSize);
    } else {
        build(builder, state, ArrayRef<int64_t>{}, type, bufferSize);
    }
}

ParseResult ChannelOp::parse(OpAsmParser &parser, OperationState &state)
{
    if (parser.parseLParen()) return failure();

    // Parse an optional buffer size
    int32_t bufferSize = 0;
    OptionalParseResult bufferSizeResult =
        parser.parseOptionalInteger(bufferSize);
    if (bufferSizeResult.has_value() && failed(*bufferSizeResult))
        return failure();
    state.addAttribute(
        getBufferSizeAttrName(state.name),
        parser.getBuilder().getI32IntegerAttr(bufferSize));

    if (parser.parseRParen() || parser.parseColon()) return failure();

    // Parse optional shape (dimension list) followed by token type
    SmallVector<int64_t> shape;
    Type tokenType;
    if (parser.parseDimensionList(shape, /*allowDynamic=*/false))
        return failure();
    if (parser.parseType(tokenType)) return failure();
    state.addAttribute(
        getTokenTypeAttrName(state.name),
        TypeAttr::get(tokenType));

    // Add result types
    state.addTypes(InputType::get(shape, tokenType));
    state.addTypes(OutputType::get(shape, tokenType));

    return success();
}

void ChannelOp::print(OpAsmPrinter &p)
{
    p << '(';
    if (const auto bufferSize = getBufferSize())
        if (bufferSize.value() != 0) p << bufferSize.value();
    p << ") : ";
    auto inputType = cast<InputType>(getInputPort().getType());
    p.printDimensionList(inputType.getShape());
    if (!inputType.getShape().empty()) p << 'x';
    p.printType(inputType.getElementType());
}

LogicalResult ChannelOp::verify()
{
    // Both ports should be used
    if (getInputPort().getUses().empty() || getOutputPort().getUses().empty())
        ::emitWarning(getLoc(), "Detect dangling channel port.");
    return success();
}

//===----------------------------------------------------------------------===//
// DFGDialect Instance Operations
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// InstantiateOp
//===----------------------------------------------------------------------===//

void InstantiateOp::build(
    OpBuilder &builder,
    OperationState &state,
    StringRef actorName,
    ValueRange inputs,
    ValueRange outputs,
    std::optional<OffloadHardware> offloaded)
{
    auto actorNameAttr = SymbolRefAttr::get(builder.getContext(), actorName);
    state.addAttribute(getActorAttrName(state.name), actorNameAttr);

    state.addOperands(inputs);
    state.addOperands(outputs);
    state.addAttribute(
        kOperandSegmentSizesAttr,
        builder.getDenseI32ArrayAttr(
            {static_cast<int32_t>(inputs.size()),
             static_cast<int32_t>(outputs.size())}));

    if (offloaded)
        state.addAttribute(
            getOffloadedAttrName(state.name),
            OffloadHardwareAttr::get(builder.getContext(), offloaded.value()));
}

void InstantiateOp::build(
    OpBuilder &builder,
    OperationState &state,
    ProcessOp actorProcess,
    ValueRange inputs,
    ValueRange outputs,
    std::optional<OffloadHardware> offloaded)
{
    build(
        builder,
        state,
        actorProcess.getNodeName(),
        inputs,
        outputs,
        offloaded);
}

void InstantiateOp::build(
    OpBuilder &builder,
    OperationState &state,
    OperatorOp actorOperator,
    ValueRange inputs,
    ValueRange outputs,
    std::optional<OffloadHardware> offloaded)
{
    build(
        builder,
        state,
        actorOperator.getNodeName(),
        inputs,
        outputs,
        offloaded);
}

ParseResult InstantiateOp::parse(OpAsmParser &parser, OperationState &state)
{
    // Optionally mark as offloaded, which optionally mark to which hardware
    if (succeeded(parser.parseOptionalKeyword("offloaded"))) {
        OffloadHardware offHW = OffloadHardware::EmitHLS; // Default backend
        if (succeeded(parser.parseEqual())) {
            StringRef hwStr;
            if (parser.parseKeyword(&hwStr)) return failure();
            auto hwVal = symbolizeOffloadHardware(hwStr);
            if (!hwVal)
                return parser.emitError(
                           parser.getNameLoc(),
                           "unknown offload hardware: ")
                       << hwStr;
            offHW = *hwVal;
        }
        state.addAttribute(
            getOffloadedAttrName(state.name),
            OffloadHardwareAttr::get(parser.getContext(), offHW));
    }

    // Parse instantiated process/operator's name
    StringAttr actorName;
    if (parser.parseSymbolName(actorName)) return failure();
    state.addAttribute(
        getActorAttrName(state.name),
        SymbolRefAttr::get(actorName));

    // Parse inputs
    SmallVector<OpAsmParser::UnresolvedOperand, 4> inputs;
    if (succeeded(parser.parseOptionalKeyword("inputs"))) {
        if (parser.parseLParen() || parser.parseOperandList(inputs) ||
            // parser.resolveOperands(inputs, opTy, result.operands)) {
            parser.parseRParen())
            return failure();
    }

    // Parse outputs
    SmallVector<OpAsmParser::UnresolvedOperand, 4> outputs;
    if (succeeded(parser.parseOptionalKeyword("outputs"))) {
        if (parser.parseLParen() || parser.parseOperandList(outputs)
            || parser.parseRParen())
            return failure();
    }

    // Parse function type
    if (parser.parseColon()) return failure();
    SMLoc location = parser.getCurrentLocation();
    FunctionType funcType;
    if (parser.parseType(funcType)) return failure();

    // Check function type and operands sizes
    auto inputTypes = funcType.getInputs();
    auto outputTypes = funcType.getResults();
    int32_t numInputs = inputTypes.size();
    int32_t numOutputs = outputTypes.size();
    if ((unsigned)numInputs != inputs.size()
        || (unsigned)numOutputs != outputs.size()) {
        parser.emitError(
            location,
            "Call signature does not match operand count");
    }

    // Resolve operands
    if (parser.resolveOperands(inputs, inputTypes, location, state.operands))
        return failure();
    if (parser.resolveOperands(outputs, outputTypes, location, state.operands))
        return failure();

    // Add derived `operand_segment_sizes` attribute based on parsed
    // operands.
    auto operandSegmentSizes =
        parser.getBuilder().getDenseI32ArrayAttr({numInputs, numOutputs});
    state.addAttribute(kOperandSegmentSizesAttr, operandSegmentSizes);

    return success();
}

void InstantiateOp::print(OpAsmPrinter &p)
{
    // Has offloaded notation
    if (auto offloaded = getOffloadedAttr())
        p << " offloaded=" << stringifyOffloadHardware(offloaded.getValue());

    // Actor name
    p << ' ';
    p.printAttributeWithoutType(getActorAttr());

    // Print inputs and output if exists
    if (!getInputs().empty()) p << " inputs(" << getInputs() << ")";
    if (!getOutputs().empty()) p << " outputs(" << getOutputs() << ")";

    // Print function type
    p << " : " << getFunctionType();
}

LogicalResult InstantiateOp::verify()
{
    auto instantiatedOp = getInstantiatedOperation();
    if (!instantiatedOp)
        return emitOpError("Cannot find embedded region operation.");

    auto baseFuncTy = getBaseFunctionType();
    auto instantiatedBaseFuncTy =
        dyn_cast<NodeInterface>(instantiatedOp).getBaseFunctionType();
    if (baseFuncTy != instantiatedBaseFuncTy)
        return emitOpError(
            "Expected same signature of embedded region and this operation.");

    return success();
}

//===----------------------------------------------------------------------===//
// EmbedOp
//===----------------------------------------------------------------------===//

void EmbedOp::build(
    OpBuilder &builder,
    OperationState &state,
    StringRef actorName,
    ValueRange inputs,
    ValueRange outputs,
    std::optional<OffloadHardware> offloaded)
{
    auto actorNameAttr = SymbolRefAttr::get(builder.getContext(), actorName);
    state.addAttribute(getActorAttrName(state.name), actorNameAttr);

    state.addOperands(inputs);
    state.addOperands(outputs);
    state.addAttribute(
        kOperandSegmentSizesAttr,
        builder.getDenseI32ArrayAttr(
            {static_cast<int32_t>(inputs.size()),
             static_cast<int32_t>(outputs.size())}));

    if (offloaded)
        state.addAttribute(
            getOffloadedAttrName(state.name),
            OffloadHardwareAttr::get(builder.getContext(), offloaded.value()));
}

void EmbedOp::build(
    OpBuilder &builder,
    OperationState &state,
    RegionOp actorRegion,
    ValueRange inputs,
    ValueRange outputs,
    std::optional<OffloadHardware> offloaded)
{
    build(
        builder,
        state,
        actorRegion.getNodeName(),
        inputs,
        outputs,
        offloaded);
}

ParseResult EmbedOp::parse(OpAsmParser &parser, OperationState &state)
{
    // Optionally mark as offloaded, which optionally mark to which hardware
    if (succeeded(parser.parseOptionalKeyword("offloaded"))) {
        OffloadHardware offHW = OffloadHardware::EmitHLS; // Default backend
        if (succeeded(parser.parseEqual())) {
            StringRef hwStr;
            if (parser.parseKeyword(&hwStr)) return failure();
            auto hwVal = symbolizeOffloadHardware(hwStr);
            if (!hwVal)
                return parser.emitError(
                           parser.getNameLoc(),
                           "unknown offload hardware: ")
                       << hwStr;
            offHW = *hwVal;
        }
        state.addAttribute(
            getOffloadedAttrName(state.name),
            OffloadHardwareAttr::get(parser.getContext(), offHW));
    }

    // Parse embed region's name
    StringAttr actorName;
    if (parser.parseSymbolName(actorName)) return failure();
    state.addAttribute(
        getActorAttrName(state.name),
        SymbolRefAttr::get(actorName));

    // Parse inputs
    SmallVector<OpAsmParser::UnresolvedOperand, 4> inputs;
    if (succeeded(parser.parseOptionalKeyword("inputs"))) {
        if (parser.parseLParen() || parser.parseOperandList(inputs) ||
            // parser.resolveOperands(inputs, opTy, result.operands)) {
            parser.parseRParen())
            return failure();
    }

    // Parse outputs
    SmallVector<OpAsmParser::UnresolvedOperand, 4> outputs;
    if (succeeded(parser.parseOptionalKeyword("outputs"))) {
        if (parser.parseLParen() || parser.parseOperandList(outputs)
            || parser.parseRParen())
            return failure();
    }

    // Parse function type
    if (parser.parseColon()) return failure();
    SMLoc location = parser.getCurrentLocation();
    FunctionType funcType;
    if (parser.parseType(funcType)) return failure();

    // Check function type and operands sizes
    auto inputTypes = funcType.getInputs();
    auto outputTypes = funcType.getResults();
    int32_t numInputs = inputTypes.size();
    int32_t numOutputs = outputTypes.size();
    if ((unsigned)numInputs != inputs.size()
        || (unsigned)numOutputs != outputs.size()) {
        parser.emitError(
            location,
            "Call signature does not match operand count");
    }

    // Resolve operands
    if (parser.resolveOperands(inputs, inputTypes, location, state.operands))
        return failure();
    if (parser.resolveOperands(outputs, outputTypes, location, state.operands))
        return failure();

    // Add derived `operand_segment_sizes` attribute based on parsed
    // operands.
    auto operandSegmentSizes =
        parser.getBuilder().getDenseI32ArrayAttr({numInputs, numOutputs});
    state.addAttribute(kOperandSegmentSizesAttr, operandSegmentSizes);

    return success();
}

void EmbedOp::print(OpAsmPrinter &p)
{
    // Has offloaded notation
    if (auto offloaded = getOffloadedAttr())
        p << " offloaded=" << stringifyOffloadHardware(offloaded.getValue());

    // Actor name
    p << ' ';
    p.printAttributeWithoutType(getActorAttr());

    // Print inputs and output if exists
    if (!getInputs().empty()) p << " inputs(" << getInputs() << ")";
    if (!getOutputs().empty()) p << " outputs(" << getOutputs() << ")";

    // Print function type
    p << " : " << getFunctionType();
}

LogicalResult EmbedOp::verify()
{
    auto embeddedOp = getEmbeddedOperation();
    if (!embeddedOp)
        return emitOpError("Cannot find embedded region operation.");

    auto baseFuncTy = getBaseFunctionType();
    auto embedBaseFuncTy =
        dyn_cast<NodeInterface>(embeddedOp).getBaseFunctionType();
    if (baseFuncTy != embedBaseFuncTy)
        return emitOpError(
            "Expected same signature of embedded region and this operation.");

    return success();
}

//===----------------------------------------------------------------------===//
// DFG structure operations.
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// OutputOp
//===----------------------------------------------------------------------===//

LogicalResult OutputOp::verify()
{
    auto operatorOutputTypes = getParentOp().getOutputPortTypes();
    auto outputTypes = getOperandTypes();

    if (outputTypes.size() != operatorOutputTypes.size())
        return emitOpError("Number of outputs must match operator's outputs.");

    for (auto [operatorOutType, thisOutType] :
         llvm::zip(operatorOutputTypes, outputTypes)) {
        if (operatorOutType != thisOutType)
            return emitOpError(
                "Output types must match operator's output types.");
    }
    return success();
}

//===----------------------------------------------------------------------===//
// PullOp
//===----------------------------------------------------------------------===//

void PullOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn)
{
    Operation* self = getOperation();
    Operation* parent = self->getParentOfType<NodeInterface>();

    int count = 0;
    parent->walk([&](PullOp op) -> WalkResult {
        if (op.getOperation() == self) return WalkResult::interrupt();
        ++count;
        return WalkResult::advance();
    });
    setNameFn(getResult(), "token" + std::to_string(count));
}

void PullOp::build(
    OpBuilder &,
    OperationState &state,
    Value readPort,
    ValueRange indices)
{
    state.addOperands(readPort);
    if (!indices.empty()) state.addOperands(indices);
    auto readPortTy = dyn_cast<OutputType>(readPort.getType()).getElementType();
    state.addTypes(readPortTy);
}

ParseResult PullOp::parse(OpAsmParser &parser, OperationState &state)
{
    // Parse read port
    OpAsmParser::UnresolvedOperand readPort;
    if (parser.parseOperand(readPort)) return failure();

    // Indices if port is shaped
    SmallVector<OpAsmParser::UnresolvedOperand, 4> indices;
    if (succeeded(parser.parseOptionalLSquare())) {
        if (parser.parseCommaSeparatedList([&]() -> ParseResult {
                OpAsmParser::UnresolvedOperand idx;
                if (parser.parseOperand(idx)) return failure();
                indices.push_back(idx);
                return success();
            })
            || parser.parseRSquare())
            return failure();
    }

    // Parse port type
    Type portType;
    if (parser.parseColon() || parser.parseType(portType)) return failure();
    state.addTypes(cast<OutputType>(portType).getElementType());

    // Resolve operands
    if (parser.resolveOperand(readPort, portType, state.operands))
        return failure();
    if (!indices.empty()) {
        for (auto idx : indices)
            if (parser.resolveOperand(
                    idx,
                    IndexType::get(parser.getContext()),
                    state.operands))
                return failure();
    }

    return success();
}

void PullOp::print(OpAsmPrinter &p)
{
    p << " " << getReadPort();
    auto indices = getIndices();
    if (!indices.empty()) {
        p << "[";
        for (size_t i = 0; i < indices.size(); ++i) {
            p << indices[i];
            if (i < indices.size() - 1) p << ", ";
        }
        p << "]";
    }
    p << " : " << getReadPort().getType();
}

LogicalResult PullOp::verify()
{
    auto portType = cast<OutputType>(getReadPort().getType());
    auto shape = portType.getShape();
    auto indices = getIndices();

    if (shape.empty() && !indices.empty())
        return emitOpError("Indices provided but read port has no shape");

    if (!shape.empty() && indices.size() != shape.size())
        return emitOpError("Expected ") << shape.size()
                                        << " index operand(s) to match the "
                                           "rank of the read port, but got "
                                        << indices.size();

    return success();
}

//===----------------------------------------------------------------------===//
// PullAsTensorOp
//===----------------------------------------------------------------------===//

void PullAsTensorOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn)
{
    Operation* self = getOperation();
    Operation* parent = self->getParentOfType<NodeInterface>();

    int count = 0;
    parent->walk([&](PullAsTensorOp op) -> WalkResult {
        if (op.getOperation() == self) return WalkResult::interrupt();
        ++count;
        return WalkResult::advance();
    });
    setNameFn(getResult(), "token_tensor" + std::to_string(count));
}

void PullAsTensorOp::build(OpBuilder &, OperationState &state, Value readPort)
{
    state.addOperands(readPort);
    auto portType = cast<OutputType>(readPort.getType());
    state.addTypes(
        RankedTensorType::get(portType.getShape(), portType.getElementType()));
}

ParseResult PullAsTensorOp::parse(OpAsmParser &parser, OperationState &state)
{
    // Parse read port
    OpAsmParser::UnresolvedOperand readPort;
    if (parser.parseOperand(readPort)) return failure();

    // Parse port type
    Type portType;
    if (parser.parseColon() || parser.parseType(portType)) return failure();
    auto portDFGType = cast<OutputType>(portType);
    state.addTypes(
        RankedTensorType::get(
            portDFGType.getShape(),
            portDFGType.getElementType()));

    // Resolve operands
    if (parser.resolveOperand(readPort, portType, state.operands))
        return failure();

    return success();
}

void PullAsTensorOp::print(OpAsmPrinter &p)
{ p << " " << getReadPort() << " : " << getReadPort().getType(); }

LogicalResult PullAsTensorOp::verify()
{
    // Check if read port is shaped
    if (getReadPort().getType().getShape().empty())
        return emitOpError("Cannot pull as tensor from a scalar read port.");
    return success();
}

//===----------------------------------------------------------------------===//
// PullAsMemRefOp
//===----------------------------------------------------------------------===//

void PullAsMemRefOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn)
{
    Operation* self = getOperation();
    Operation* parent = self->getParentOfType<NodeInterface>();

    int count = 0;
    parent->walk([&](PullAsMemRefOp op) -> WalkResult {
        if (op.getOperation() == self) return WalkResult::interrupt();
        ++count;
        return WalkResult::advance();
    });
    setNameFn(getResult(), "token_memref" + std::to_string(count));
}

void PullAsMemRefOp::build(OpBuilder &, OperationState &state, Value readPort)
{
    state.addOperands(readPort);
    auto portType = cast<OutputType>(readPort.getType());
    state.addTypes(
        MemRefType::get(portType.getShape(), portType.getElementType()));
}

ParseResult PullAsMemRefOp::parse(OpAsmParser &parser, OperationState &state)
{
    // Parse read port
    OpAsmParser::UnresolvedOperand readPort;
    if (parser.parseOperand(readPort)) return failure();

    // Parse port type
    Type portType;
    if (parser.parseColon() || parser.parseType(portType)) return failure();
    auto portDFGType = cast<OutputType>(portType);
    state.addTypes(
        MemRefType::get(portDFGType.getShape(), portDFGType.getElementType()));

    // Resolve operands
    if (parser.resolveOperand(readPort, portType, state.operands))
        return failure();

    return success();
}

void PullAsMemRefOp::print(OpAsmPrinter &p)
{ p << " " << getReadPort() << " : " << getReadPort().getType(); }

LogicalResult PullAsMemRefOp::verify()
{
    // Check if read port is shaped
    if (getReadPort().getType().getShape().empty())
        return emitOpError("Cannot pull as tensor from a scalar read port.");
    return success();
}

//===----------------------------------------------------------------------===//
// PushOp
//===----------------------------------------------------------------------===//

void PushOp::build(
    OpBuilder &,
    OperationState &state,
    Value token,
    Value writePort,
    ValueRange indices)
{
    state.addOperands(token);
    state.addOperands(writePort);
    if (!indices.empty()) state.addOperands(indices);
}

ParseResult PushOp::parse(OpAsmParser &parser, OperationState &state)
{
    // Parse token
    OpAsmParser::UnresolvedOperand token;
    if (parser.parseOperand(token)) return failure();

    // Parse write port
    OpAsmParser::UnresolvedOperand writePort;
    if (parser.parseKeyword("to") || parser.parseOperand(writePort))
        return failure();

    // Indices if port is shaped
    SmallVector<OpAsmParser::UnresolvedOperand, 4> indices;
    if (succeeded(parser.parseOptionalLSquare())) {
        if (parser.parseCommaSeparatedList([&]() -> ParseResult {
                OpAsmParser::UnresolvedOperand idx;
                if (parser.parseOperand(idx)) return failure();
                indices.push_back(idx);
                return success();
            })
            || parser.parseRSquare())
            return failure();
    }

    // Parse port type
    Type portType;
    if (parser.parseColon() || parser.parseType(portType)) return failure();

    // Resolve operands
    if (parser.resolveOperand(
            token,
            cast<InputType>(portType).getElementType(),
            state.operands))
        return failure();
    if (parser.resolveOperand(writePort, portType, state.operands))
        return failure();
    if (!indices.empty()) {
        for (auto idx : indices)
            if (parser.resolveOperand(
                    idx,
                    IndexType::get(parser.getContext()),
                    state.operands))
                return failure();
    }

    return success();
}

void PushOp::print(OpAsmPrinter &p)
{
    p << " " << getToken() << " to " << getWritePort();
    auto indices = getIndices();
    if (!indices.empty()) {
        p << "[";
        for (size_t i = 0; i < indices.size(); ++i) {
            p << indices[i];
            if (i < indices.size() - 1) p << ", ";
        }
        p << "]";
    }
    p << " : " << getWritePort().getType();
}

LogicalResult PushOp::verify()
{
    auto portType = cast<InputType>(getWritePort().getType());
    auto shape = portType.getShape();
    auto indices = getIndices();

    if (shape.empty() && !indices.empty())
        return emitOpError("Indices provided but read port has no shape");

    if (!shape.empty() && indices.size() != shape.size())
        return emitOpError("Expected ") << shape.size()
                                        << " index operand(s) to match the "
                                           "rank of the read port, but got "
                                        << indices.size();

    return success();
}

//===----------------------------------------------------------------------===//
// PushTensorOp
//===----------------------------------------------------------------------===//

void PushTensorOp::build(
    OpBuilder &,
    OperationState &state,
    Value tokenTensor,
    Value writePort)
{
    state.addOperands(tokenTensor);
    state.addOperands(writePort);
}

ParseResult PushTensorOp::parse(OpAsmParser &parser, OperationState &state)
{
    // Parse token
    OpAsmParser::UnresolvedOperand token;
    if (parser.parseOperand(token)) return failure();

    // Parse write port
    OpAsmParser::UnresolvedOperand writePort;
    if (parser.parseKeyword("to") || parser.parseOperand(writePort))
        return failure();

    // Parse port type
    Type portType;
    if (parser.parseColon() || parser.parseType(portType)) return failure();

    // Resolve operands
    auto portDFGType = cast<InputType>(portType);
    if (parser.resolveOperand(
            token,
            RankedTensorType::get(
                portDFGType.getShape(),
                portDFGType.getElementType()),
            state.operands))
        return failure();
    if (parser.resolveOperand(writePort, portType, state.operands))
        return failure();

    return success();
}

void PushTensorOp::print(OpAsmPrinter &p)
{
    p << " " << getTokenTensor() << " to " << getWritePort() << " : "
      << getWritePort().getType();
}

//===----------------------------------------------------------------------===//
// PushMemRefOp
//===----------------------------------------------------------------------===//

void PushMemRefOp::build(
    OpBuilder &,
    OperationState &state,
    Value tokenMemref,
    Value writePort)
{
    state.addOperands(tokenMemref);
    state.addOperands(writePort);
}

ParseResult PushMemRefOp::parse(OpAsmParser &parser, OperationState &state)
{
    // Parse token
    OpAsmParser::UnresolvedOperand token;
    if (parser.parseOperand(token)) return failure();

    // Parse write port
    OpAsmParser::UnresolvedOperand writePort;
    if (parser.parseKeyword("to") || parser.parseOperand(writePort))
        return failure();

    // Parse port type
    Type portType;
    if (parser.parseColon() || parser.parseType(portType)) return failure();

    // Resolve operands
    auto portDFGType = cast<InputType>(portType);
    if (parser.resolveOperand(
            token,
            MemRefType::get(
                portDFGType.getShape(),
                portDFGType.getElementType()),
            state.operands))
        return failure();
    if (parser.resolveOperand(writePort, portType, state.operands))
        return failure();

    return success();
}

void PushMemRefOp::print(OpAsmPrinter &p)
{
    p << " " << getTokenMemref() << " to " << getWritePort() << " : "
      << getWritePort().getType();
}

//===----------------------------------------------------------------------===//
// DFGDialect
//===----------------------------------------------------------------------===//

void DFGDialect::registerOps()
{
    addOperations<
#define GET_OP_LIST
#include "laksa-mlir/Dialect/DFG/IR/DFGOps.cpp.inc"
        >();
}