Program Listing for File EmitHLSOps.cpp

↰ Return to documentation for file (lib/Dialect/EmitHLS/IR/EmitHLSOps.cpp)

#include "laksa-mlir/Dialect/EmitHLS/IR/EmitHLSOps.h"

#include "laksa-mlir/Dialect/EmitHLS/EmitHLSEnums.h"
#include "laksa-mlir/Dialect/EmitHLS/IR/EmitHLSBase.h"
#include "laksa-mlir/Dialect/EmitHLS/IR/EmitHLSTypes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/IRMapping.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 <cstdint>
#include <llvm/ADT/STLExtras.h>
#include <llvm/ADT/SmallVector.h>
#include <llvm/ADT/StringRef.h>
#include <llvm/Support/LogicalResult.h>
#include <mlir/IR/BuiltinAttributes.h>
#include <mlir/IR/BuiltinOps.h>
#include <mlir/IR/MLIRContext.h>
#include <mlir/IR/OpImplementation.h>
#include <mlir/IR/OperationSupport.h>
#include <mlir/IR/PatternMatch.h>
#include <mlir/IR/ValueRange.h>
#include <mlir/Support/LLVM.h>

using namespace mlir;
using namespace mlir::emithls;

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

#define GET_OP_CLASSES
#include "laksa-mlir/Dialect/EmitHLS/IR/EmitHLSOps.cpp.inc"

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

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

//===----------------------------------------------------------------------===//
// Structure operations
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// IncludeOp
//===----------------------------------------------------------------------===//

void IncludeOp::build(
    OpBuilder &builder,
    OperationState &state,
    StringRef header)
{
    state.addAttribute(
        getHeaderAttrName(state.name),
        builder.getStringAttr(header));
}

ParseResult IncludeOp::parse(OpAsmParser &parser, OperationState &state)
{
    StringAttr includeStr;
    OptionalParseResult includeParseResult = parser.parseOptionalAttribute(
        includeStr,
        getHeaderAttrName(state.name),
        state.attributes);
    if (!includeParseResult.has_value())
        return parser.emitError(parser.getNameLoc())
               << "expected string attribute";

    return success();
}

void IncludeOp::print(OpAsmPrinter &p) { p << " \"" << getHeader() << "\""; }

LogicalResult IncludeOp::verify()
{
    auto parent = getOperation()->getParentOp();
    if (!isa<ModuleOp>(parent))
        return emitOpError(
            "An IncludeOp must only have ModuleOp as its parent.");
    return success();
}

namespace {
struct RemoveSameInclude final : public OpRewritePattern<IncludeOp> {
    using OpRewritePattern<IncludeOp>::OpRewritePattern;
    LogicalResult
    matchAndRewrite(IncludeOp op, PatternRewriter &rewriter) const override
    {
        auto moduleOp = op->getParentOfType<ModuleOp>();
        for (auto &bodyOp : moduleOp.getBody()->getOperations()) {
            if (&bodyOp == op.getOperation()) return failure();
            if (auto prevInclude = dyn_cast<IncludeOp>(&bodyOp))
                if (prevInclude.getHeader() == op.getHeader()) {
                    rewriter.eraseOp(op);
                    return success();
                }
        }
        return failure();
    }
};
struct HoistIncludeToTop final : public OpRewritePattern<IncludeOp> {
    using OpRewritePattern<IncludeOp>::OpRewritePattern;
    LogicalResult
    matchAndRewrite(IncludeOp op, PatternRewriter &rewriter) const override
    {
        auto moduleOp = op->getParentOfType<ModuleOp>();
        Block* moduleBody = moduleOp.getBody();

        Operation* firstNonInclude = nullptr;
        for (auto &opInModule : moduleBody->getOperations()) {
            if (!isa<IncludeOp>(&opInModule)) {
                firstNonInclude = &opInModule;
                break;
            }
        }

        // If there is no other operation
        if (!firstNonInclude) return failure();

        for (auto &opInModule : moduleBody->getOperations()) {
            if (&opInModule == op.getOperation()) return failure();
            if (&opInModule == firstNonInclude) break;
        }

        rewriter.moveOpBefore(op, firstNonInclude);
        return success();
    }
};
struct ReorderIncludes final : public OpRewritePattern<IncludeOp> {
    using OpRewritePattern<IncludeOp>::OpRewritePattern;
    LogicalResult
    matchAndRewrite(IncludeOp op, PatternRewriter &rewriter) const override
    {
        auto moduleOp = op->getParentOfType<ModuleOp>();

        SmallVector<IncludeOp> includes(
            moduleOp.getBody()->getOps<IncludeOp>());

        if (includes.size() <= 1) return failure();

        auto currentIdx =
            std::distance(includes.begin(), llvm::find(includes, op));
        auto sortedIdx = llvm::count_if(includes, [&](IncludeOp includeOp) {
            return includeOp != op && includeOp.getHeader() < op.getHeader();
        });

        if (currentIdx == sortedIdx)
            return failure();
        else if (currentIdx < sortedIdx)
            rewriter.moveOpAfter(op, includes[sortedIdx]);
        else
            rewriter.moveOpBefore(op, includes[sortedIdx]);

        return success();
    }
};
} // namespace

void IncludeOp::getCanonicalizationPatterns(
    RewritePatternSet &patterns,
    MLIRContext* context)
{
    patterns.add<RemoveSameInclude>(context);
    patterns.add<HoistIncludeToTop>(context);
    patterns.add<ReorderIncludes>(context);
}

//===----------------------------------------------------------------------===//
// FuncOp
//===----------------------------------------------------------------------===//

void FuncOp::build(
    OpBuilder &builder,
    OperationState &state,
    StringRef name,
    FunctionType functionType)
{ build(builder, state, name, functionType, ArrayAttr{}, ArrayAttr{}); }

ParseResult FuncOp::parse(OpAsmParser &parser, OperationState &state)
{
    auto buildFuncType = [](Builder &builder,
                            ArrayRef<Type> argTypes,
                            ArrayRef<Type> results,
                            function_interface_impl::VariadicFlag,
                            std::string &) {
        return builder.getFunctionType(argTypes, results);
    };

    return function_interface_impl::parseFunctionOp(
        parser,
        state,
        /*allowVariadic=*/false,
        getFunctionTypeAttrName(state.name),
        buildFuncType,
        getArgAttrsAttrName(state.name),
        getResAttrsAttrName(state.name));
}

void FuncOp::print(OpAsmPrinter &p)
{
    function_interface_impl::printFunctionOp(
        p,
        *this,
        /*isVariadic=*/false,
        getFunctionTypeAttrName(),
        getArgAttrsAttrName(),
        getResAttrsAttrName());
}

LogicalResult FuncOp::verify()
{
    auto funcTy = getFunctionType();
    if (funcTy.getNumInputs() != 0) {
        for (auto type : funcTy.getInputs())
            if (auto arrayType = dyn_cast<ArrayType>(type)) {
                auto arrayElemType = arrayType.getElementType();
                if (isa<PointerType>(arrayElemType))
                    return emitOpError("Array of pointer type is not allowed.");
            } else if (!isa<StreamType, PointerType>(type)) {
                return emitOpError("Unsupported type ")
                       << type
                       << ", needs to be stream/pointer/array type for "
                          "function arguments.";
            }
    }
    return success();
}

//===----------------------------------------------------------------------===//
// CallOp
//===----------------------------------------------------------------------===//

void CallOp::build(
    OpBuilder &,
    OperationState &state,
    FuncOp callee,
    ValueRange operands)
{
    state.addOperands(operands);
    state.addAttribute(
        getCalleeAttrName(state.name),
        SymbolRefAttr::get(callee));
    state.addTypes(callee.getFunctionType().getResults());
}

//===----------------------------------------------------------------------===//
// VariableOp
//===----------------------------------------------------------------------===//

void VariableOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn)
{
    auto thisOp = getOperation();

    // Count sibling variables within the nearest enclosing FuncOp when there
    // is one; otherwise fall back to the nearest ancestor op that sits
    // directly under the module.
    Operation* scopeOp = thisOp->getParentOfType<FuncOp>();
    if (!scopeOp) {
        scopeOp = thisOp;
        while (Operation* parent = scopeOp->getParentOp()) {
            if (isa<ModuleOp>(parent)) break;
            scopeOp = parent;
        }
        if (scopeOp == thisOp || scopeOp->getNumRegions() == 0) return;
    }

    auto resultTy = getType();
    unsigned count = 0;
    for (auto &opi : scopeOp->getRegion(0).front()) {
        if (auto varOp = dyn_cast<VariableOp>(opi)) {
            if (varOp.getOperation() == thisOp) break;
            if (((varOp.getType() == resultTy)
                 || (isa<ArrayType>(varOp.getType())
                     && isa<ArrayType>(resultTy)))
                && (varOp.getIsConst() == getIsConst()))
                ++count;
        }
    }

    std::string prefix = getIsConst() ? "const" : "var";
    std::string suffix;
    if (isa<IntegerType>(resultTy))
        suffix =
            "_int" + std::to_string(resultTy.getIntOrFloatBitWidth()) + "_";
    else if (isa<IndexType>(resultTy))
        suffix = "_index_";
    else if (isa<ArrayType>(resultTy))
        suffix = "_array_";
    else if (isa<StreamType>(resultTy))
        suffix = "_stream_";
    setNameFn(getVariable(), prefix + suffix + std::to_string(count));
}

void VariableOp::build(
    OpBuilder &builder,
    OperationState &state,
    Type type,
    Attribute initNumber,
    bool is_const)
{
    state.addTypes(type);
    if (initNumber)
        state.addAttribute(getInitNumberAttrName(state.name), initNumber);
    if (is_const) state.addAttribute("is_const", builder.getUnitAttr());
}

void VariableOp::build(
    OpBuilder &builder,
    OperationState &state,
    Type type,
    Value initValue,
    bool is_const)
{
    state.addTypes(type);
    if (initValue) state.addOperands(initValue);
    if (is_const) state.addAttribute("is_const", builder.getUnitAttr());
}

ParseResult VariableOp::parse(OpAsmParser &parser, OperationState &state)
{
    if (failed(parser.parseKeyword("as")))
        return parser.emitError(
            parser.getCurrentLocation(),
            "expected keyword 'as'");

    // If it's const value
    bool isConst = false;
    if (succeeded(parser.parseOptionalKeyword("const"))) {
        state.addAttribute("is_const", parser.getBuilder().getUnitAttr());
        isConst = true;
    }

    // Parse type
    Type type;
    if (failed(parser.parseType(type)))
        return parser.emitError(parser.getCurrentLocation(), "expected type");
    state.addTypes(type);

    Attribute initNumber;
    OpAsmParser::UnresolvedOperand initValue;

    if (succeeded(parser.parseOptionalEqual())) {
        auto operandParseResult = parser.parseOptionalOperand(initValue);
        if (operandParseResult.has_value()) {
            // If there is an initial value
            if (failed(*operandParseResult)) return failure();
            if (failed(parser.resolveOperand(initValue, type, state.operands)))
                return failure();
        } else if (
            failed(parser.parseAttribute(
                initNumber,
                type,
                getInitNumberAttrName(state.name).data(),
                state.attributes))) {
            // If there is an initial attribute
            return parser.emitError(
                parser.getCurrentLocation(),
                "expected initial attribute");
        }
    } else {
        if (isConst)
            return parser.emitError(
                parser.getCurrentLocation(),
                "const value must have initial attribute");
    }

    if (failed(parser.parseOptionalAttrDict(state.attributes)))
        return failure();

    return success();
}

void VariableOp::print(OpAsmPrinter &p)
{
    p << " as ";
    if (getIsConst()) p << "const ";
    p << getType();
    if (getInitNumber()) {
        p << " = ";
        p.printAttributeWithoutType(getInitNumberAttr());
    } else if (getInitValue()) {
        p << " = " << getInitValue();
    }
    p.printOptionalAttrDict(
        (*this)->getAttrs(),
        {getIsConstAttrName(), getInitNumberAttrName()});
}

LogicalResult VariableOp::verify()
{
    auto resultType = getType();
    if (isa<PointerType>(resultType))
        return emitOpError("Dynamic memory usage is not supported in HLS.");
    if (getInitNumber()
        && !isa<IntegerType, IndexType, FloatType, ArrayType>(resultType))
        return emitOpError("Unsupported type to have init value.");
    return success();
}

namespace {
struct RemoveUnusedVariable final : public OpRewritePattern<VariableOp> {
    using OpRewritePattern<VariableOp>::OpRewritePattern;

    LogicalResult
    matchAndRewrite(VariableOp op, PatternRewriter &rewriter) const override
    {
        if (op.getVariable().use_empty()) {
            rewriter.eraseOp(op);
            return success();
        }
        return failure();
    }
};

struct RemoveSameConstVariable final : public OpRewritePattern<VariableOp> {
    using OpRewritePattern<VariableOp>::OpRewritePattern;

    LogicalResult
    matchAndRewrite(VariableOp op, PatternRewriter &rewriter) const override
    {
        if (!op.getIsConst()) return failure();

        auto initAttr = op.getInitNumberAttr();
        auto constTy = op.getType();

        // Dedupe within the nearest enclosing FuncOp when there is one;
        // otherwise fall back to the nearest ancestor op that sits directly
        // under the module.
        Operation* scopeOp = op->getParentOfType<FuncOp>();
        if (!scopeOp) {
            scopeOp = op;
            while (Operation* parent = scopeOp->getParentOp()) {
                if (isa<ModuleOp>(parent)) break;
                scopeOp = parent;
            }
            if (scopeOp == op.getOperation()) return failure();
        }

        VariableOp candidate;
        scopeOp->walk([&](VariableOp varOp) -> WalkResult {
            if (varOp == op) return WalkResult::interrupt();
            if (varOp.getIsConst() && varOp->getParentOp() == scopeOp
                && varOp.getInitNumberAttr() == initAttr
                && varOp.getType() == constTy)
                candidate = varOp;
            return WalkResult::advance();
        });

        if (!candidate) return failure();
        op.getResult().replaceAllUsesWith(candidate.getResult());
        rewriter.eraseOp(op);
        return success();
    }
};

struct HoistToFunctionTop final : public OpRewritePattern<VariableOp> {
    using OpRewritePattern<VariableOp>::OpRewritePattern;

    LogicalResult
    matchAndRewrite(VariableOp op, PatternRewriter &) const override
    {
        if (!op.getIsConst()) return failure();

        if (auto parentFunc = op->getParentOfType<FuncOp>()) {
            if (isa<FuncOp>(op->getParentOp())) return failure();
            auto &entryBlock = parentFunc.getRegion().front();
            op->moveBefore(&entryBlock, entryBlock.begin());
            return success();
        }

        // Hoist to the top of the nearest ancestor op that sits directly under
        // the module
        Operation* topLevelOp = op;
        while (Operation* parent = topLevelOp->getParentOp()) {
            if (isa<ModuleOp>(parent)) break;
            topLevelOp = parent;
        }
        if (topLevelOp == op.getOperation() || topLevelOp->getNumRegions() == 0)
            return failure();
        if (op->getParentOp() == topLevelOp) return failure();

        auto &entryBlock = topLevelOp->getRegion(0).front();
        op->moveBefore(&entryBlock, entryBlock.begin());
        return success();
    }
};
} // namespace

void VariableOp::getCanonicalizationPatterns(
    RewritePatternSet &patterns,
    MLIRContext* context)
{
    patterns.add<RemoveUnusedVariable>(context);
    patterns.add<RemoveSameConstVariable>(context);
    patterns.add<HoistToFunctionTop>(context);
}

//===----------------------------------------------------------------------===//
// UpdateOp
//===----------------------------------------------------------------------===//

void UpdateOp::build(
    OpBuilder &builder,
    OperationState &state,
    Value variable,
    Value newValue,
    SmallVector<Value> indices,
    SmallVector<Value> newValueIndices)
{
    state.addOperands(variable);
    if (!indices.empty()) state.addOperands(indices);
    state.addOperands(newValue);
    if (!newValueIndices.empty()) state.addOperands(newValueIndices);
    auto operandSegmentSizes =
        builder.getDenseI32ArrayAttr({/*variable*/ 1,
                                      int32_t(indices.size()),
                                      /*newValue*/ 1,
                                      int32_t(newValueIndices.size())});
    state.addAttribute(kOperandSegmentSizesAttr, operandSegmentSizes);
}

ParseResult UpdateOp::parse(OpAsmParser &parser, OperationState &state)
{
    OpAsmParser::UnresolvedOperand variable, newValue;
    SmallVector<OpAsmParser::UnresolvedOperand, 4> indices, newValueIndices;

    if (parser.parseOperand(variable)) return failure();
    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();
    }

    if (parser.parseKeyword("with") || parser.parseOperand(newValue))
        return failure();
    if (succeeded(parser.parseOptionalLSquare())) {
        if (parser.parseCommaSeparatedList([&]() -> ParseResult {
                OpAsmParser::UnresolvedOperand idx;
                if (parser.parseOperand(idx)) return failure();
                newValueIndices.push_back(idx);
                return success();
            })
            || parser.parseRSquare())
            return failure();
    }

    Type variableTy, newValueTy;
    if (parser.parseColon() || parser.parseType(variableTy)
        || parser.parseLess() || parser.parseMinus()
        || parser.parseType(newValueTy))
        return failure();

    if (parser.resolveOperand(variable, variableTy, state.operands))
        return failure();
    if (!indices.empty()) {
        for (auto idx : indices)
            if (parser.resolveOperand(
                    idx,
                    parser.getBuilder().getIndexType(),
                    state.operands))
                return failure();
    }
    if (parser.resolveOperand(newValue, newValueTy, state.operands))
        return failure();
    if (!newValueIndices.empty()) {
        for (auto idx : newValueIndices)
            if (parser.resolveOperand(
                    idx,
                    parser.getBuilder().getIndexType(),
                    state.operands))
                return failure();
    }

    auto operandSegmentSizes = parser.getBuilder().getDenseI32ArrayAttr(
        {/*variable*/ 1,
         int32_t(indices.size()),
         /*newValue*/ 1,
         int32_t(newValueIndices.size())});
    state.addAttribute(kOperandSegmentSizesAttr, operandSegmentSizes);

    return success();
}

void UpdateOp::print(OpAsmPrinter &p)
{
    p << " " << getVariable();
    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 << " with " << getNewValue();
    auto newValueIndices = getNewValueIndices();
    if (!newValueIndices.empty()) {
        p << "[";
        for (size_t i = 0; i < newValueIndices.size(); ++i) {
            p << newValueIndices[i];
            if (i < newValueIndices.size() - 1) p << ", ";
        }
        p << "]";
    }
    p << " : " << getVariable().getType() << " <- " << getNewValue().getType();
}

LogicalResult UpdateOp::verify()
{
    auto variable = getVariable();
    Type varTy =
        getIndices().empty()
            ? variable.getType()
            : dyn_cast<ShapedType>(variable.getType()).getElementType();
    auto newValue = getNewValue();
    Type newValTy =
        getNewValueIndices().empty()
            ? newValue.getType()
            : dyn_cast<ShapedType>(newValue.getType()).getElementType();
    if (varTy != newValTy)
        return emitOpError("update value and variable types don't match");
    return success();
}

//===----------------------------------------------------------------------===//
// ExpressionOp
//===----------------------------------------------------------------------===//

void ExpressionOp::build(
    OpBuilder &builder,
    OperationState &state,
    Type resultType,
    function_ref<void(OpBuilder &, Location)> bodyBuilder)
{
    OpBuilder::InsertionGuard guard(builder);
    Region* bodyRegion = state.addRegion();
    builder.createBlock(bodyRegion);
    bodyBuilder(builder, state.location);
    state.addTypes(resultType);
}

void ExpressionOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn)
{
    Operation* self = getOperation();
    Operation* topLevel = self;
    while (Operation* p = topLevel->getParentOp()) {
        if (isa<ModuleOp>(p)) break;
        topLevel = p;
    }

    unsigned count = 0;
    topLevel->walk([&](ExpressionOp op) -> WalkResult {
        if (op.getOperation() == self) return WalkResult::interrupt();
        ++count;
        return WalkResult::advance();
    });
    setNameFn(getResult(), "expr" + std::to_string(count));
}

LogicalResult ExpressionOp::verify()
{
    auto yieldOp = cast<YieldOp>((&getBody().front())->getTerminator());
    if (getType() != yieldOp.getValue().getType())
        return emitOpError(
            "Expression return a different type from yield value.");
    return success();
}

namespace {
struct InlinePureNestedExpression final
        : public OpRewritePattern<ExpressionOp> {
    using OpRewritePattern<ExpressionOp>::OpRewritePattern;

    LogicalResult
    matchAndRewrite(ExpressionOp op, PatternRewriter &rewriter) const override
    {
        Block* block = &op.getBody().front();

        // Match: body must be exactly an ExpressionOp followed by a YieldOp.
        if (std::ranges::distance(block->getOperations()) != 2)
            return failure();
        auto innerExpr = dyn_cast<ExpressionOp>(block->front());
        if (!innerExpr) return failure();

        // Splice the inner expression's body (including its YieldOp) before it,
        // then erase the now-redundant wrapper and the outer YieldOp.
        rewriter.setInsertionPointToStart(block);
        IRMapping mapper;
        for (auto &opi : innerExpr.getBody().front())
            rewriter.clone(opi, mapper);
        rewriter.eraseOp(&block->back()); // outer YieldOp
        rewriter.eraseOp(innerExpr);
        return success();
    }
};
} // namespace

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

//===----------------------------------------------------------------------===//
// YieldOp
//===----------------------------------------------------------------------===//

void YieldOp::build(OpBuilder &, OperationState &state, Value value)
{ state.addOperands(value); }

void YieldOp::build(OpBuilder &, OperationState &)
{
    // Do nothing
}

//===----------------------------------------------------------------------===//
// IfOp
//===----------------------------------------------------------------------===//

void IfOp::build(
    OpBuilder &builder,
    OperationState &state,
    Value condition,
    function_ref<void(OpBuilder &, Location)> thenBuilder,
    function_ref<void(OpBuilder &, Location)> elseBuilder)
{
    state.addOperands(condition);

    OpBuilder::InsertionGuard guard(builder);
    Region* thenRegion = state.addRegion();
    builder.createBlock(thenRegion);
    thenBuilder(builder, state.location);
    Region* elseRegion = state.addRegion();
    if (elseBuilder) {
        builder.createBlock(elseRegion);
        elseBuilder(builder, state.location);
    }
}

void IfOp::build(
    OpBuilder &builder,
    OperationState &state,
    Value condition,
    bool ensureElseBlock)
{
    state.addOperands(condition);

    OpBuilder::InsertionGuard guard(builder);
    Region* thenRegion = state.addRegion();
    builder.createBlock(thenRegion);
    Region* elseRegion = state.addRegion();
    if (ensureElseBlock) builder.createBlock(elseRegion);
}

ParseResult IfOp::parse(OpAsmParser &parser, OperationState &state)
{
    state.regions.reserve(2);
    Region* thenRegion = state.addRegion();
    Region* elseRegion = state.addRegion();

    OpAsmParser::UnresolvedOperand condition;
    if (parser.parseOperand(condition)
        || parser.resolveOperand(
            condition,
            parser.getBuilder().getI1Type(),
            state.operands))
        return failure();

    if (parser.parseRegion(*thenRegion, {}, {})) return failure();
    if (succeeded(parser.parseOptionalKeyword("else")))
        if (parser.parseRegion(*elseRegion, {}, {})) return failure();

    return success();
}

void IfOp::print(OpAsmPrinter &p)
{
    p << " " << getCondition() << " ";
    p.printRegion(getThenRegion(), false, false);
    auto &elseRegion = getElseRegion();
    if (!elseRegion.empty()) {
        p << " else ";
        p.printRegion(elseRegion, false, false);
    }
}

//===----------------------------------------------------------------------===//
// ForOp
//===----------------------------------------------------------------------===//

void ForOp::build(
    OpBuilder &builder,
    OperationState &state,
    int64_t lowerBound,
    int64_t upperBound,
    int64_t step,
    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder)
{
    state.addAttribute(
        getLowerBoundAttrName(state.name),
        builder.getIntegerAttr(builder.getIndexType(), lowerBound));
    state.addAttribute(
        getUpperBoundAttrName(state.name),
        builder.getIntegerAttr(builder.getIndexType(), upperBound));
    state.addAttribute(
        getStepAttrName(state.name),
        builder.getIntegerAttr(builder.getIndexType(), step));

    OpBuilder::InsertionGuard guard(builder);
    Region* bodyRegion = state.addRegion();
    Block* bodyBlock = builder.createBlock(bodyRegion);
    bodyBlock->addArgument(builder.getIndexType(), state.location);
    if (bodyBuilder)
        bodyBuilder(builder, state.location, bodyBlock->getArguments());
}

void ForOp::build(
    OpBuilder &builder,
    OperationState &state,
    Attribute lowerBound,
    Attribute upperBound,
    Attribute step,
    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder)
{
    state.addAttribute(getLowerBoundAttrName(state.name), lowerBound);
    state.addAttribute(getUpperBoundAttrName(state.name), upperBound);
    state.addAttribute(getStepAttrName(state.name), step);

    OpBuilder::InsertionGuard g(builder);
    Region* bodyRegion = state.addRegion();
    Block* bodyBlock = builder.createBlock(bodyRegion);
    bodyBlock->addArgument(builder.getIndexType(), state.location);
    if (bodyBuilder)
        bodyBuilder(builder, state.location, bodyBlock->getArguments());
}

void ForOp::getAsmBlockArgumentNames(
    Region &region,
    OpAsmSetValueNameFn setNameFn)
{
    Operation* forOp = getOperation();
    unsigned nestedLevel = 0;

    while ((forOp = forOp->getParentOp()) && forOp)
        if (isa<ForOp>(forOp)) ++nestedLevel;

    setNameFn(region.getArgument(0), "idx" + std::to_string(nestedLevel));
}

ParseResult ForOp::parse(OpAsmParser &parser, OperationState &state)
{
    auto &builder = parser.getBuilder();
    Type indexType = builder.getIndexType();

    OpAsmParser::Argument inductionVariable;
    if (parser.parseOperand(inductionVariable.ssaName) || parser.parseEqual())
        return failure();

    IntegerAttr lbAttr, ubAttr, stepAttr;
    if (parser.parseAttribute(
            lbAttr,
            indexType,
            getLowerBoundAttrName(state.name).data(),
            state.attributes)
        || parser.parseKeyword("to")
        || parser.parseAttribute(
            ubAttr,
            indexType,
            getUpperBoundAttrName(state.name).data(),
            state.attributes)
        || parser.parseKeyword("step")
        || parser.parseAttribute(
            stepAttr,
            indexType,
            getStepAttrName(state.name).data(),
            state.attributes))
        return failure();

    SmallVector<OpAsmParser::Argument> blkArgs;
    blkArgs.push_back(inductionVariable);
    blkArgs.front().type = indexType;

    Region* body = state.addRegion();
    if (parser.parseRegion(*body, blkArgs)) return failure();

    if (parser.parseOptionalAttrDict(state.attributes)) return failure();

    return success();
}

void ForOp::print(OpAsmPrinter &p)
{
    p << " " << getInductionVariable() << " = " << getLowerBound() << " to "
      << getUpperBound() << " step " << getStep();
    p << ' ';
    p.printRegion(
        getRegion(),
        /*printEntryBlockArgs=*/false,
        /*printBlockTerminators=*/false);
    p.printOptionalAttrDict(
        (*this)->getAttrs(),
        {getLowerBoundAttrName(), getUpperBoundAttrName(), getStepAttrName()});
}

LogicalResult ForOp::verify()
{
    if (getLowerBound() == getUpperBound())
        return emitOpError("expected loop iterate at least once");
    if (getStep() == 0) return emitOpError("expected non-zero step");
    if (getUpperBound().slt(getLowerBound() + getStep()))
        return emitOpError("unexpected index overflow");
    return success();
}

namespace {
struct RemoveTrivialLoop : public OpRewritePattern<ForOp> {
    using OpRewritePattern<ForOp>::OpRewritePattern;

    LogicalResult
    matchAndRewrite(ForOp op, PatternRewriter &rewriter) const override
    {
        if ((op.getLowerBound() + op.getStep()) != op.getUpperBound())
            return failure();

        Block &body = op.getRegion().front();
        bool hasPipeline = !body.empty() && isa<PragmaPipelineOp>(body.front());
        auto parent = dyn_cast<ForOp>(op->getParentOp());
        bool siblingHasLoop =
            parent
            && llvm::any_of(parent.getBody().front(), [&](Operation &sibling) {
                   return &sibling != op.getOperation()
                          && sibling
                                 .walk([](ForOp) {
                                     return WalkResult::interrupt();
                                 })
                                 .wasInterrupted();
               });
        if (hasPipeline && (!parent || siblingHasLoop)) return failure();

        if (hasPipeline)
            rewriter.moveOpBefore(
                &body.front(),
                &parent.getBody().front(),
                parent.getBody().front().begin());

        rewriter.setInsertionPoint(op);
        IRMapping mapper;
        auto constVarOp = VariableOp::create(
            rewriter,
            op.getLoc(),
            rewriter.getIndexType(),
            op.getLowerBoundAttr(),
            true);
        mapper.map(op.getInductionVariable(), constVarOp.getResult());
        for (auto &opFor : op.getRegion().getOps())
            rewriter.clone(opFor, mapper);
        rewriter.eraseOp(op);
        return success();
    }
};
struct RemoveEmptyLoop : public OpRewritePattern<ForOp> {
    using OpRewritePattern<ForOp>::OpRewritePattern;

    LogicalResult
    matchAndRewrite(ForOp op, PatternRewriter &rewriter) const override
    {
        // Preserve loops whose body is not empty
        if (!op.getRegion().getOps().empty()) return failure();
        rewriter.eraseOp(op);
        return success();
    }
};
struct CollapseUnusedOuterLoop : public OpRewritePattern<ForOp> {
    using OpRewritePattern<ForOp>::OpRewritePattern;

    LogicalResult
    matchAndRewrite(ForOp op, PatternRewriter &rewriter) const override
    {
        // Only safe when NEITHER induction variable is read
        if (!op.getInductionVariable().use_empty()) return failure();

        Block &body = op.getRegion().front();
        if (body.empty() || &body.front() != &body.back()) return failure();
        auto inner = dyn_cast<ForOp>(body.front());
        if (!inner || !inner.getInductionVariable().use_empty())
            return failure();

        rewriter.setInsertionPoint(op);
        auto merged = ForOp::create(
            rewriter,
            op.getLoc(),
            0,
            op.getTripCount() * inner.getTripCount());
        // Move inner's body into merged
        rewriter.mergeBlocks(
            &inner.getRegion().front(),
            &merged.getRegion().front(),
            merged.getInductionVariable());
        rewriter.eraseOp(op);
        return success();
    }
};
} // namespace

void ForOp::getCanonicalizationPatterns(
    RewritePatternSet &results,
    MLIRContext* context)
{ results.add<RemoveTrivialLoop, CollapseUnusedOuterLoop>(context); }

//===----------------------------------------------------------------------===//
// Arithmetic operations
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// ArithFusedOp
//===----------------------------------------------------------------------===//

LogicalResult ArithFusedOp::verify()
{
    if (auto varOp = dyn_cast<VariableOp>(getAcc().getDefiningOp())) {
        if (varOp.getIsConst())
            return emitOpError("Accumulator cannot be a constant variable.");
    } else {
        return emitOpError("Accumulator of the fused op must be a variable.");
    }
    return success();
}

//===----------------------------------------------------------------------===//
// ArithDataRangeOp
//===----------------------------------------------------------------------===//

ParseResult ArithDataRangeOp::parse(OpAsmParser &parser, OperationState &state)
{
    OpAsmParser::UnresolvedOperand data, hBit, lBit;
    Type dataTy, outTy;
    Type indexTy = parser.getBuilder().getIndexType();
    if (parser.parseOperand(data) || parser.parseLParen()
        || parser.parseOperand(hBit) || parser.parseComma()
        || parser.parseOperand(lBit) || parser.parseRParen())
        return failure();
    if (parser.parseColon() || parser.parseType(dataTy) || parser.parseArrow()
        || parser.parseType(outTy))
        return failure();
    if (parser.resolveOperand(data, dataTy, state.operands)
        || parser.resolveOperand(hBit, indexTy, state.operands)
        || parser.resolveOperand(lBit, indexTy, state.operands))
        return failure();
    state.addTypes(outTy);
    return success();
}

void ArithDataRangeOp::print(OpAsmPrinter &p)
{
    p << " " << getData() << "(" << getHighBit() << ", " << getLowBit() << ")";
    p << " : " << getData().getType() << " -> " << getType();
}

//===----------------------------------------------------------------------===//
// Array operations
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// ArrayReadOp
//===----------------------------------------------------------------------===//

LogicalResult ArrayReadOp::verify()
{
    if (static_cast<int64_t>(getIndices().size()) != getArrayType().getRank()) {
        return emitOpError(
                   "incorrect number of indices for array read, expected ")
               << getArrayType().getRank() << " but got "
               << getIndices().size();
    }
    return success();
}

//===----------------------------------------------------------------------===//
// ArrayWriteOp
//===----------------------------------------------------------------------===//

LogicalResult ArrayWriteOp::verify()
{
    if (static_cast<int64_t>(getIndices().size()) != getArrayType().getRank())
        return emitOpError(
            "array write index operand count not equal to memref rank");
    return success();
}

//===----------------------------------------------------------------------===//
// Stream operations
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// StreamReadOp
//===----------------------------------------------------------------------===//

void StreamReadOp::build(
    OpBuilder &,
    OperationState &state,
    Type resultType,
    Value stream,
    ValueRange indices)
{
    state.addOperands(stream);
    if (!indices.empty()) state.addOperands(indices);
    state.addTypes(resultType);
}

LogicalResult StreamReadOp::verify()
{
    auto streamType = getStream().getType();
    auto indices = getIndices();
    Type elemType;

    if (auto arrayType = dyn_cast<ArrayType>(streamType)) {
        // If the stream is an array of streams
        auto arrayShape = arrayType.getShape();
        auto arrayElemType = arrayType.getElementType();
        if (!isa<StreamType>(arrayElemType))
            return emitOpError("expect array of streams");
        elemType = cast<StreamType>(arrayElemType).getElementType();
        // Then, check the indices
        if (indices.empty()) return emitOpError("expect indices");
        if (indices.size() != arrayShape.size())
            return emitOpError("indices and array shapes mismatch");
    } else {
        // Otherwise, it must be a stream type
        elemType = cast<StreamType>(streamType).getElementType();
        // Expect no indices
        if (!indices.empty()) return emitOpError("expect no indices");
    }

    if (elemType != getType()) emitOpError("result and stream types mismatch");

    return success();
}

//===----------------------------------------------------------------------===//
// StreamWriteOp
//===----------------------------------------------------------------------===//

void StreamWriteOp::build(
    OpBuilder &,
    OperationState &state,
    Value data,
    Value stream,
    ValueRange indices)
{
    state.addOperands(data);
    state.addOperands(stream);
    if (!indices.empty()) state.addOperands(indices);
}

LogicalResult StreamWriteOp::verify()
{
    auto streamType = getStream().getType();
    auto indices = getIndices();
    Type elemType;

    if (auto arrayType = dyn_cast<ArrayType>(streamType)) {
        // If the stream is an array of streams
        auto arrayShape = arrayType.getShape();
        auto arrayElemType = arrayType.getElementType();
        if (!isa<StreamType>(arrayElemType))
            return emitOpError("expect array of streams");
        elemType = cast<StreamType>(arrayElemType).getElementType();
        // Then, check the indices
        if (indices.empty()) return emitOpError("expect indices");
        if (indices.size() != arrayShape.size())
            return emitOpError("indices and array shapes mismatch");
    } else {
        // Otherwise, it must be a stream type
        elemType = cast<StreamType>(streamType).getElementType();
        // Expect no indices
        if (!indices.empty()) return emitOpError("expect no indices");
    }

    if (elemType != getData().getType())
        return emitOpError("data and stream types mismatch");
    return success();
}

//===----------------------------------------------------------------------===//
// Pragma operations
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// PragmaArrayPartitionOp
//===----------------------------------------------------------------------===//

void PragmaArrayPartitionOp::build(
    OpBuilder &builder,
    OperationState &state,
    Value variable,
    ArrayPartitionType partType,
    std::optional<int32_t> partFactor,
    std::optional<int32_t> partDim)
{
    state.addOperands(variable);
    state.addAttribute(
        getPartTypeAttrName(state.name),
        ArrayPartitionTypeAttr::get(builder.getContext(), partType));
    if (partFactor.has_value()) {
        state.addAttribute(
            getPartFactorAttrName(state.name),
            builder.getI32IntegerAttr(partFactor.value()));
    }
    if (partDim.has_value()) {
        state.addAttribute(
            getPartDimAttrName(state.name),
            builder.getI32IntegerAttr(partDim.value()));
    }
}

ParseResult
PragmaArrayPartitionOp::parse(OpAsmParser &parser, OperationState &state)
{
    OpAsmParser::UnresolvedOperand variable;
    Type variableType;
    if (parser.parseKeyword("variable") || parser.parseEqual()
        || parser.parseOperand(variable) || parser.parseLParen()
        || parser.parseType(variableType) || parser.parseRParen())
        return failure();
    if (parser.resolveOperand(variable, variableType, state.operands))
        return failure();

    StringRef partTypeStr;
    if (parser.parseKeyword("type") || parser.parseEqual()
        || parser.parseKeyword(&partTypeStr))
        return failure();
    auto partType = symbolizeArrayPartitionType(partTypeStr);
    if (!partType.has_value())
        return parser.emitError(
            parser.getNameLoc(),
            "unknown array partition type");
    state.addAttribute(
        getPartTypeAttrName(state.name),
        ArrayPartitionTypeAttr::get(parser.getContext(), partType.value()));

    if (succeeded(parser.parseOptionalKeyword("factor"))) {
        int32_t partFactor;
        if (parser.parseEqual() || parser.parseInteger(partFactor))
            return failure();
        state.addAttribute(
            getPartFactorAttrName(state.name),
            parser.getBuilder().getI32IntegerAttr(partFactor));
    }

    if (succeeded(parser.parseOptionalKeyword("dim"))) {
        int32_t partDim;
        if (parser.parseEqual() || parser.parseInteger(partDim))
            return failure();
        state.addAttribute(
            getPartDimAttrName(state.name),
            parser.getBuilder().getI32IntegerAttr(partDim));
    }

    return success();
}

void PragmaArrayPartitionOp::print(OpAsmPrinter &p)
{
    p << " variable=" << getVariable() << "(" << getVariableType() << ")";
    p << " type=" << stringifyArrayPartitionType(getPartType());
    if (getPartFactor()) p << " factor=" << getPartFactor();
    if (getPartDim()) p << " dim=" << getPartDim();
}

LogicalResult PragmaArrayPartitionOp::verify()
{
    auto partType = getPartType();
    auto partFactor = getPartFactor();
    auto partDim = getPartDim();

    if (partType == ArrayPartitionType::cyclic
        || partType == ArrayPartitionType::block) {
        if (!partFactor.has_value())
            return emitOpError("cyclic and block partition requires a factor");
        if (!partDim.has_value())
            return emitOpError("cyclic and block partition requires a dim");
    } else {
        if (partFactor.has_value())
            return emitOpError("complete partition needs no factor");
    }
    return success();
}

//===----------------------------------------------------------------------===//
// PragmaBindStorageOp
//===----------------------------------------------------------------------===//

void PragmaBindStorageOp::build(
    OpBuilder &builder,
    OperationState &state,
    Value variable,
    BindStorageType storageType,
    BindStorageImpl storageImpl)
{
    state.addOperands(variable);
    state.addAttribute(
        getStorageTypeAttrName(state.name),
        BindStorageTypeAttr::get(builder.getContext(), storageType));
    state.addAttribute(
        getStorageImplAttrName(state.name),
        BindStorageImplAttr::get(builder.getContext(), storageImpl));
}

ParseResult
PragmaBindStorageOp::parse(OpAsmParser &parser, OperationState &state)
{
    OpAsmParser::UnresolvedOperand variable;
    Type variableType;
    if (parser.parseKeyword("variable") || parser.parseEqual()
        || parser.parseOperand(variable) || parser.parseLParen()
        || parser.parseType(variableType) || parser.parseRParen())
        return failure();
    if (parser.resolveOperand(variable, variableType, state.operands))
        return failure();

    StringRef storageTypeStr;
    if (parser.parseKeyword("type") || parser.parseEqual()
        || parser.parseKeyword(&storageTypeStr))
        return failure();
    auto storageType = symbolizeBindStorageType(storageTypeStr);
    if (!storageType.has_value())
        return parser.emitError(
            parser.getNameLoc(),
            "unknown bind storage type");
    state.addAttribute(
        getStorageTypeAttrName(state.name),
        BindStorageTypeAttr::get(
            parser.getBuilder().getContext(),
            storageType.value()));

    StringRef storageImplStr;
    if (parser.parseKeyword("impl") || parser.parseEqual()
        || parser.parseKeyword(&storageImplStr))
        return failure();
    auto storageImpl = symbolizeBindStorageImpl(storageImplStr);
    if (!storageImpl.has_value())
        return parser.emitError(
            parser.getNameLoc(),
            "unknown bind storage impl");
    state.addAttribute(
        getStorageImplAttrName(state.name),
        BindStorageImplAttr::get(
            parser.getBuilder().getContext(),
            storageImpl.value()));

    return success();
}

void PragmaBindStorageOp::print(OpAsmPrinter &p)
{
    p << " variable=" << getVariable() << "(" << getVariableType() << ")";
    p << " type=" << stringifyBindStorageType(getStorageType());
    p << " impl=" << stringifyBindStorageImpl(getStorageImpl());
}

//===----------------------------------------------------------------------===//
// PragmaDataflowOp
//===----------------------------------------------------------------------===//

void PragmaDataflowOp::build(
    OpBuilder &builder,
    OperationState &state,
    function_ref<void(OpBuilder &, Location)> bodyBuilder)
{
    OpBuilder::InsertionGuard guard(builder);
    Region* dataflowRegion = state.addRegion();
    if (bodyBuilder) {
        builder.createBlock(dataflowRegion);
        bodyBuilder(builder, state.location);
    }
}

//===----------------------------------------------------------------------===//
// PragmaInlineOp
//===----------------------------------------------------------------------===//

void PragmaInlineOp::build(OpBuilder &builder, OperationState &state, bool off)
{
    if (off)
        state.addAttribute(getOffAttrName(state.name), builder.getUnitAttr());
}

ParseResult PragmaInlineOp::parse(OpAsmParser &parser, OperationState &state)
{
    if (succeeded(parser.parseOptionalKeyword("off")))
        state.addAttribute(
            getOffAttrName(state.name),
            parser.getBuilder().getUnitAttr());
    return success();
}

void PragmaInlineOp::print(OpAsmPrinter &p)
{
    if (getOff()) p << " off";
}

namespace {
struct HoistPragmaInlinePattern : public OpRewritePattern<PragmaInlineOp> {
    using OpRewritePattern<PragmaInlineOp>::OpRewritePattern;

    LogicalResult
    matchAndRewrite(PragmaInlineOp op, PatternRewriter &) const override
    {
        Block* block = op->getBlock();
        if (&block->front() == op.getOperation()) return failure();
        op->moveBefore(&block->front());
        return success();
    }
    // Ensure this has higher priority over constant hoist
    PatternBenefit getBenefit() const { return 2; }
};
} // namespace

void PragmaInlineOp::getCanonicalizationPatterns(
    RewritePatternSet &patterns,
    MLIRContext* context)
{ patterns.add<HoistPragmaInlinePattern>(context, 100); }

//===----------------------------------------------------------------------===//
// PragmaPipelineOp
//===----------------------------------------------------------------------===//

void PragmaPipelineOp::build(
    OpBuilder &builder,
    OperationState &result,
    int32_t interval,
    emithls::PipelineStyle style)
{
    result.addAttribute(
        getIntervalAttrName(result.name),
        builder.getI32IntegerAttr(interval));
    result.addAttribute(
        getStyleAttrName(result.name),
        emithls::PipelineStyleAttr::get(builder.getContext(), style));
}

ParseResult PragmaPipelineOp::parse(OpAsmParser &parser, OperationState &state)
{
    // Parse initial interval as integer
    int32_t interval;
    if (parser.parseKeyword("II") || parser.parseEqual()
        || parser.parseInteger(interval))
        return failure();
    state.addAttribute(
        getIntervalAttrName(state.name),
        parser.getBuilder().getI32IntegerAttr(interval));

    // Parse style string, if none, use flp by default
    PipelineStyle style = PipelineStyle::flp;
    if (succeeded(parser.parseOptionalKeyword("style"))) {
        if (parser.parseEqual()) return failure();
        StringRef styleStr;
        if (parser.parseKeyword(&styleStr)) return failure();
        auto styleValue = symbolizePipelineStyle(styleStr);
        if (!styleValue)
            return parser.emitError(
                parser.getNameLoc(),
                "unknown pipeline style");
        style = *styleValue;
    }
    state.addAttribute(
        getStyleAttrName(state.name),
        PipelineStyleAttr::get(parser.getContext(), style));

    return success();
}

void PragmaPipelineOp::print(OpAsmPrinter &p)
{
    p << " II=" << getInterval()
      << " style=" << stringifyPipelineStyle(getStyle());
}

LogicalResult PragmaPipelineOp::verify()
{
    if (!isa<ForOp>(getOperation()->getParentOp()))
        return emitOpError("pipeline pragma must be in a for-loop");
    if (getInterval() <= 0) return emitOpError("II must be greater than 0");
    return success();
}

namespace {
struct RemoveNestedPipeline final : public OpRewritePattern<PragmaPipelineOp> {
    using OpRewritePattern<PragmaPipelineOp>::OpRewritePattern;

    LogicalResult matchAndRewrite(
        PragmaPipelineOp op,
        PatternRewriter &rewriter) const override
    {
        auto hasPrevPipeline = [](Operation* o) {
            auto* prev = o->getPrevNode();
            return prev && isa<PragmaPipelineOp>(prev);
        };
        // Check if there is pipeline before
        if (hasPrevPipeline(op)) {
            rewriter.eraseOp(op);
            return success();
        }
        // Check if the parent loop is already pipelined
        for (Operation* parent = op->getParentOp(); parent;
             parent = parent->getParentOp()) {
            if (isa<emithls::ForOp>(parent) && hasPrevPipeline(parent)) {
                rewriter.eraseOp(op);
                return success();
            }
        }
        return failure();
    }
};

} // namespace

void PragmaPipelineOp::getCanonicalizationPatterns(
    RewritePatternSet &results,
    MLIRContext* context)
{ results.add<RemoveNestedPipeline>(context); }

//===----------------------------------------------------------------------===//
// PragmaStreamOp
//===----------------------------------------------------------------------===//

void PragmaStreamOp::build(
    OpBuilder &builder,
    OperationState &state,
    Value variable,
    int32_t depth)
{
    state.addOperands(variable);
    state.addAttribute(
        getDepthAttrName(state.name),
        builder.getI32IntegerAttr(depth));
}

ParseResult PragmaStreamOp::parse(OpAsmParser &parser, OperationState &state)
{
    OpAsmParser::UnresolvedOperand variable;
    Type variableType;
    if (parser.parseKeyword("variable") || parser.parseEqual()
        || parser.parseOperand(variable) || parser.parseLParen()
        || parser.parseType(variableType) || parser.parseRParen())
        return failure();
    if (parser.resolveOperand(variable, variableType, state.operands))
        return failure();

    int32_t depth;
    if (parser.parseKeyword("depth") || parser.parseEqual()
        || parser.parseInteger(depth))
        return failure();
    state.addAttribute(
        getDepthAttrName(state.name),
        parser.getBuilder().getI32IntegerAttr(depth));

    return success();
}

void PragmaStreamOp::print(OpAsmPrinter &p)
{
    p << " variable=" << getVariable() << "(" << getVariableType() << ")";
    p << " depth=" << getDepth();
}

LogicalResult PragmaStreamOp::verify()
{
    if (getDepth() <= 0) return emitOpError("expect positive integer depth");
    if (auto arrayType = dyn_cast<ArrayType>(getVariableType())) {
        if (!isa<StreamType>(arrayType.getElementType()))
            return emitOpError("expect array of streams");
    }
    return success();
}

//===----------------------------------------------------------------------===//
// PragmaUnrollOp
//===----------------------------------------------------------------------===//

void PragmaUnrollOp::build(
    OpBuilder &builder,
    OperationState &state,
    std::optional<int32_t> factor)
{
    if (factor.has_value())
        state.addAttribute(
            getFactorAttrName(state.name),
            builder.getI32IntegerAttr(factor.value()));
}

ParseResult PragmaUnrollOp::parse(OpAsmParser &parser, OperationState &state)
{
    if (succeeded(parser.parseOptionalKeyword("factor"))) {
        int32_t factor;
        if (parser.parseEqual() || parser.parseInteger(factor))
            return failure();
        state.addAttribute(
            getFactorAttrName(state.name),
            parser.getBuilder().getI32IntegerAttr(factor));
    }
    return success();
}

void PragmaUnrollOp::print(OpAsmPrinter &p)
{
    if (getFactor()) p << " factor=" << getFactor();
}

//===----------------------------------------------------------------------===//
// Helper Operations
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// HelperLineBufferOp
//===----------------------------------------------------------------------===//

void HelperLineBufferOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn)
{
    Operation* self = getOperation();
    Operation* topLevel = self;
    while (Operation* p = topLevel->getParentOp()) {
        if (isa<ModuleOp>(p)) break;
        topLevel = p;
    }
    unsigned count = 0;
    topLevel->walk([&](HelperLineBufferOp op) -> WalkResult {
        if (op.getOperation() == self) return WalkResult::interrupt();
        ++count;
        return WalkResult::advance();
    });
    setNameFn(getBufRef(), "linebuf_" + std::to_string(count));
}

void HelperLineBufferOp::build(
    OpBuilder &builder,
    OperationState &state,
    Value tokenRef,
    Value index,
    Type linebufType,
    ArrayRef<int32_t> numChan,
    int32_t numLine)
{
    state.addOperands(tokenRef);
    if (index) state.addOperands(index);
    state.addAttribute(
        getNumChanAttrName(state.name),
        builder.getI32ArrayAttr(numChan));
    state.addAttribute(
        getNumLineAttrName(state.name),
        builder.getI32IntegerAttr(numLine));
    state.addTypes({linebufType});
}

ParseResult
HelperLineBufferOp::parse(OpAsmParser &parser, OperationState &state)
{
    OpAsmParser::UnresolvedOperand tokenRef, index;
    bool hasKeepIndex = false;
    int32_t numLine;
    SmallVector<int32_t> numChan;
    Type tokenType, linebufType;

    if (parser.parseOperand(tokenRef) || parser.parseKeyword("num_chan")
        || parser.parseLSquare())
        return failure();

    if (parser.parseOptionalRSquare().failed()) {
        do {
            int32_t val;
            if (parser.parseInteger(val)) return failure();
            numChan.push_back(val);
        } while (parser.parseOptionalComma().succeeded());
        if (parser.parseRSquare()) return failure();
    }

    if (parser.parseKeyword("num_line") || parser.parseInteger(numLine))
        return failure();
    if (succeeded(parser.parseOptionalKeyword("keep"))) {
        hasKeepIndex = true;
        if (parser.parseOperand(index)) return failure();
    }
    if (parser.parseColon() || parser.parseType(tokenType)
        || parser.parseArrow() || parser.parseType(linebufType))
        return failure();

    if (parser.resolveOperand(tokenRef, tokenType, state.operands))
        return failure();
    if (hasKeepIndex)
        if (parser.resolveOperand(
                index,
                parser.getBuilder().getIndexType(),
                state.operands))
            return failure();

    state.addAttribute(
        getNumChanAttrName(state.name),
        parser.getBuilder().getI32ArrayAttr(numChan));
    state.addAttribute(
        getNumLineAttrName(state.name),
        parser.getBuilder().getI32IntegerAttr(numLine));
    state.addTypes({linebufType});

    return success();
}

void HelperLineBufferOp::print(OpAsmPrinter &p)
{
    p << " " << getTokenRef() << " num_chan [";
    llvm::interleaveComma(getNumChan(), p, [&](Attribute attr) {
        p << llvm::cast<IntegerAttr>(attr).getInt();
    });
    p << "] num_line " << getNumLine();
    if (getIndex()) p << " keep " << getIndex();
    p << " : " << getTokenRef().getType() << " -> " << getBufRef().getType();
}

//===----------------------------------------------------------------------===//
// HelperWindowOp
//===----------------------------------------------------------------------===//

void HelperWindowOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn)
{
    Operation* self = getOperation();
    Operation* topLevel = self;
    while (Operation* p = topLevel->getParentOp()) {
        if (isa<ModuleOp>(p)) break;
        topLevel = p;
    }
    unsigned count = 0;
    topLevel->walk([&](HelperWindowOp op) -> WalkResult {
        if (op.getOperation() == self) return WalkResult::interrupt();
        ++count;
        return WalkResult::advance();
    });
    setNameFn(getWinRef(), "window_" + std::to_string(count));
}

void HelperWindowOp::build(
    OpBuilder &,
    OperationState &state,
    Value linebuf,
    ValueRange indices,
    Type windowType)
{
    state.addOperands(linebuf);
    state.addOperands(indices);
    state.addTypes(windowType);
}

ParseResult HelperWindowOp::parse(OpAsmParser &parser, OperationState &state)
{
    OpAsmParser::UnresolvedOperand linebuf;
    SmallVector<OpAsmParser::UnresolvedOperand> indices;
    Type bufferType, windowType;

    if (parser.parseOperand(linebuf) || parser.parseKeyword("from")
        || parser.parseLSquare()
        || parser.parseCommaSeparatedList([&]() -> ParseResult {
               OpAsmParser::UnresolvedOperand idx;
               if (parser.parseOperand(idx)) return failure();
               indices.push_back(idx);
               return success();
           })
        || parser.parseRSquare() || parser.parseColon()
        || parser.parseType(bufferType) || parser.parseArrow()
        || parser.parseType(windowType))
        return failure();

    if (parser.resolveOperand(linebuf, bufferType, state.operands))
        return failure();
    if (parser.resolveOperands(
            indices,
            parser.getBuilder().getIndexType(),
            state.operands))
        return failure();
    state.addTypes(windowType);

    return success();
}

void HelperWindowOp::print(OpAsmPrinter &p)
{
    p << " " << getBufRef() << " from [";
    for (auto [i, index] : llvm::enumerate(getIndices())) {
        p << index;
        if (i == getIndices().size() - 1) break;
        p << ", ";
    }
    p << "] : " << getBufRef().getType() << " -> " << getWinRef().getType();
}

//===----------------------------------------------------------------------===//
// HelperAccumulateOp
//===----------------------------------------------------------------------===//

void HelperAccumulateOp::build(
    OpBuilder &builder,
    OperationState &state,
    Value accuRef,
    ValueRange indices,
    FusedOperator operatorCode,
    Value value)
{
    state.addOperands(accuRef);
    state.addOperands(value);
    state.addOperands(indices);
    state.addAttribute(
        getOpCodeAttrName(state.name),
        FusedOperatorAttr::get(builder.getContext(), operatorCode));
}

ParseResult
HelperAccumulateOp::parse(OpAsmParser &parser, OperationState &state)
{
    OpAsmParser::UnresolvedOperand accuRef, value;
    SmallVector<OpAsmParser::UnresolvedOperand> indices;
    StringRef operatorCodeStr;
    Type type;

    if (parser.parseOperand(accuRef) || parser.parseKeyword("at")
        || parser.parseLSquare())
        return failure();
    if (failed(parser.parseOptionalRSquare())) {
        if (parser.parseCommaSeparatedList([&]() -> ParseResult {
                OpAsmParser::UnresolvedOperand idx;
                if (parser.parseOperand(idx)) return failure();
                indices.push_back(idx);
                return success();
            })
            || parser.parseRSquare())
            return failure();
    }
    if (parser.parseKeyword(&operatorCodeStr) || parser.parseOperand(value)
        || parser.parseColonType(type))
        return failure();

    auto operatorCode = symbolizeFusedOperator(operatorCodeStr);
    if (!operatorCode.has_value())
        return parser.emitError(
            parser.getNameLoc(),
            "unknown accumulate operator");
    state.addAttribute(
        getOpCodeAttrName(state.name),
        FusedOperatorAttr::get(parser.getContext(), operatorCode.value()));

    auto shapedType = cast<ShapedType>(type);
    Type valueType = shapedType.getElementType();
    if (succeeded(parser.parseOptionalKeyword("cast"))
        && parser.parseType(valueType))
        return failure();

    if (parser.resolveOperand(accuRef, shapedType, state.operands))
        return failure();
    if (parser.resolveOperand(value, valueType, state.operands))
        return failure();
    if (parser.resolveOperands(
            indices,
            parser.getBuilder().getIndexType(),
            parser.getCurrentLocation(),
            state.operands))
        return failure();

    return success();
}

void HelperAccumulateOp::print(OpAsmPrinter &p)
{
    p << " " << getAccuRef() << " at [";
    llvm::interleaveComma(getIndices(), p);
    p << "] " << stringifyFusedOperator(getOpCode()) << " " << getValue();
    p << " : " << getAccuRef().getType();
    Type valueType = getValue().getType();
    if (valueType != cast<ShapedType>(getAccuRef().getType()).getElementType())
        p << " cast " << valueType;
}

//===----------------------------------------------------------------------===//
// EmitHLSDialect
//===----------------------------------------------------------------------===//

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