From db9ce6456d1b4ae593b60e04f81d318760fdf0ae Mon Sep 17 00:00:00 2001 From: Hanchen Ye Date: Fri, 29 Mar 2024 22:56:44 -0500 Subject: [PATCH] Replace itensor/tensor_init op with itensor/tensor_instance op after dataflow scheduling --- include/scalehls/Dialect/HLS/IR/HLSOps.td | 26 ++++++---- .../BufferizableOpInterfaceImpl.cpp | 47 +++++++------------ .../HLS/Transforms/LowerITensorToStream.cpp | 26 ++++++---- .../HLS/Transforms/ScalarizeITensor.cpp | 21 +++++---- .../HLS/Transforms/ScheduleDataflow.cpp | 30 +++++++++++- 5 files changed, 90 insertions(+), 60 deletions(-) diff --git a/include/scalehls/Dialect/HLS/IR/HLSOps.td b/include/scalehls/Dialect/HLS/IR/HLSOps.td index d2f87c3e..444754f2 100644 --- a/include/scalehls/Dialect/HLS/IR/HLSOps.td +++ b/include/scalehls/Dialect/HLS/IR/HLSOps.td @@ -343,11 +343,18 @@ def StreamOp : HLSOp<"stream", [ DeclareOpInterfaceMethods]> { let summary = "Instantiate a stream channel"; + let arguments = (ins OptionalAttr:$location); let results = (outs AnyStream:$stream); - let assemblyFormat = "attr-dict `:` type(results)"; + let assemblyFormat = [{ + (`location` $location^)? attr-dict `:` type(results) + }]; let hasVerifier = 1; let hasCanonicalizeMethod = 1; + let builders = [ + OpBuilder<(ins "mlir::Type":$stream), + "build($_builder, $_state, stream, nullptr);"> + ]; let extraClassDeclaration = [{ SmallVector getReadUses(); @@ -411,7 +418,7 @@ def BufferOp : HLSOp<"buffer", [ OptionalAttr:$location); let results = (outs AnyBuffer:$memref); let assemblyFormat = [{ - (`init` $initValue^)? attr-dict `:` type(results) + (`init` $initValue^)? (`location` $location^)? attr-dict `:` type(results) }]; let hasVerifier = 1; @@ -435,6 +442,13 @@ def TaskOp : HLSOp<"task", [SingleBlockImplicitTerminator<"YieldOp">]> { OptionalAttr:$location); let results = (outs Variadic:$results); let regions = (region SizedRegion<1>:$body); + let assemblyFormat = [{ + $name (`location` $location^)? (`inits` $inits^)? $body attr-dict `:` + functional-type(operands, results) + }]; + + let hasVerifier = 1; + let hasCanonicalizer = 1; let builders = [ OpBuilder<(ins "mlir::ValueRange":$inits, "mlir::StringAttr":$name, "mlir::StringAttr":$location), @@ -445,14 +459,6 @@ def TaskOp : HLSOp<"task", [SingleBlockImplicitTerminator<"YieldOp">]> { "build($_builder, $_state, inits, nullptr);"> ]; - let assemblyFormat = [{ - $name (`location` $location^)? (`inits` $inits^)? $body attr-dict `:` - functional-type(operands, results) - }]; - - let hasVerifier = 1; - let hasCanonicalizer = 1; - let extraClassDeclaration = [{ TypeRange getInitTypes() { return TypeRange(getInits()); } diff --git a/lib/Dialect/HLS/Transforms/BufferizableOpInterfaceImpl.cpp b/lib/Dialect/HLS/Transforms/BufferizableOpInterfaceImpl.cpp index 09fbba9b..0efc950a 100644 --- a/lib/Dialect/HLS/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/lib/Dialect/HLS/Transforms/BufferizableOpInterfaceImpl.cpp @@ -448,9 +448,9 @@ struct YieldOpInterface } }; -struct TensorInitOpInterface - : public BufferizableOpInterface::ExternalModel { +struct TensorInstanceOpInterface + : public BufferizableOpInterface::ExternalModel { bool bufferizesToAllocation(Operation *op, Value value) const { return true; } bool resultBufferizesToMemoryWrite(Operation *op, OpResult opResult, @@ -462,53 +462,40 @@ struct TensorInitOpInterface LogicalResult bufferize(Operation *op, RewriterBase &rewriter, const BufferizationOptions &options) const { OpBuilder::InsertionGuard g(rewriter); - auto tensorInit = cast(op); + auto tensorInst = cast(op); - // Nothing to do for dead TensorInitOps. - if (tensorInit->getUses().empty()) { - rewriter.eraseOp(tensorInit); + // Nothing to do for dead TensorInstanceOps. + if (tensorInst->getUses().empty()) { + rewriter.eraseOp(tensorInst); return success(); } - // Create memory allocation. auto maybeType = - bufferization::getBufferType(tensorInit.getResult(), options); + bufferization::getBufferType(tensorInst.getResult(), options); if (failed(maybeType)) return failure(); - for (auto &use : llvm::make_early_inc_range(tensorInit->getUses())) { - rewriter.setInsertionPoint(use.getOwner()); - FailureOr buffer = options.createAlloc( - rewriter, tensorInit.getLoc(), maybeType->cast(), {}); - if (failed(buffer)) - return failure(); - - // Handle initial value. - auto bufferOp = buffer->getDefiningOp(); - bufferOp.setInitValueAttr(tensorInit.getInitValueAttr()); - - auto repl = rewriter.create( - tensorInit.getLoc(), *buffer); - rewriter.updateRootInPlace(use.getOwner(), [&]() { use.set(repl); }); - } - rewriter.eraseOp(tensorInit); + auto buffer = rewriter.create( + tensorInst.getLoc(), *maybeType, tensorInst.getInitValueAttr(), + tensorInst.getLocationAttr()); + replaceOpWithBufferizedValues(rewriter, op, buffer.getResult()); return success(); } FailureOr getBufferType(Operation *op, Value value, const BufferizationOptions &options, SmallVector &invocationStack) const { - auto tensorInit = cast(op); - assert(value == tensorInit.getResult() && "invalid value"); + auto tensorInst = cast(op); + assert(value == tensorInst.getResult() && "invalid value"); // Compute memory space of this allocation. Attribute memorySpace; if (options.defaultMemorySpace.has_value()) memorySpace = *options.defaultMemorySpace; else - return tensorInit.emitError("could not infer memory space"); + return tensorInst.emitError("could not infer memory space"); - return getMemRefTypeWithStaticIdentityLayout(tensorInit.getType(), + return getMemRefTypeWithStaticIdentityLayout(tensorInst.getType(), memorySpace); } }; @@ -518,6 +505,6 @@ void mlir::scalehls::hls::registerBufferizableOpInterfaceExternalModels( registry.addExtension(+[](MLIRContext *ctx, HLSDialect *dialect) { hls::TaskOp::attachInterface(*ctx); hls::YieldOp::attachInterface(*ctx); - hls::TensorInitOp::attachInterface(*ctx); + hls::TensorInstanceOp::attachInterface(*ctx); }); } diff --git a/lib/Dialect/HLS/Transforms/LowerITensorToStream.cpp b/lib/Dialect/HLS/Transforms/LowerITensorToStream.cpp index d014baf0..704ed12e 100644 --- a/lib/Dialect/HLS/Transforms/LowerITensorToStream.cpp +++ b/lib/Dialect/HLS/Transforms/LowerITensorToStream.cpp @@ -25,6 +25,23 @@ static StreamType getStreamType(hls::ITensorType iTensorType) { iTensorType.getDepth()); } +namespace { +struct LowerITensorInstanceOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(hls::ITensorInstanceOp inst, + PatternRewriter &rewriter) const override { + auto streamType = getStreamType(inst.getType()); + auto stream = rewriter.create(inst.getLoc(), streamType, + inst.getLocationAttr()); + rewriter.replaceOpWithNewOp(inst, inst.getType(), + stream); + return success(); + } +}; +} // namespace + namespace { struct LowerITensorReadOp : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -144,14 +161,6 @@ struct RemoveITensorToStreamOp if (toStream.use_empty()) { rewriter.eraseOp(toStream); return success(); - } - - if (auto init = toStream.getITensor().getDefiningOp()) { - if (!init.getInitValue()) { - rewriter.replaceOpWithNewOp(toStream, - toStream.getStreamType()); - return success(); - } } else if (auto toITensor = toStream.getITensor() .getDefiningOp()) { if (toStream.getStreamType() == toITensor.getStreamType()) { @@ -217,6 +226,7 @@ struct LowerITensorToStream // Apply lowering patterns. mlir::RewritePatternSet patterns(context); + patterns.add(context); patterns.add(context); patterns.add(context); patterns.add(context); diff --git a/lib/Dialect/HLS/Transforms/ScalarizeITensor.cpp b/lib/Dialect/HLS/Transforms/ScalarizeITensor.cpp index 60fac102..a177c0cc 100644 --- a/lib/Dialect/HLS/Transforms/ScalarizeITensor.cpp +++ b/lib/Dialect/HLS/Transforms/ScalarizeITensor.cpp @@ -50,22 +50,23 @@ static ITensorType getScalarITensorType(ITensorType iTensor) { } namespace { -struct ScalarizeITensorInitOp : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; +struct ScalarizeITensorInstanceOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(hls::ITensorInitOp init, + LogicalResult matchAndRewrite(hls::ITensorInstanceOp inst, PatternRewriter &rewriter) const override { - auto iTensorType = init.getType(); + auto iTensorType = inst.getType(); if (!iTensorType.hasShapedElementType()) return failure(); - rewriter.updateRootInPlace(init, [&]() { - init.getResult().setType(getScalarITensorType(iTensorType)); + rewriter.updateRootInPlace(inst, [&]() { + inst.getResult().setType(getScalarITensorType(iTensorType)); }); - rewriter.setInsertionPointAfter(init); + rewriter.setInsertionPointAfter(inst); auto cast = - rewriter.create(init.getLoc(), iTensorType, init); - rewriter.replaceAllUsesExcept(init, cast.getResult(), cast); + rewriter.create(inst.getLoc(), iTensorType, inst); + rewriter.replaceAllUsesExcept(inst, cast.getResult(), cast); return success(); } }; @@ -279,7 +280,7 @@ struct ScalarizeITensor // Apply scalarization patterns. mlir::RewritePatternSet patterns(context); - patterns.add(context); + patterns.add(context); patterns.add(context); patterns.add(context); patterns.add(context); diff --git a/lib/Dialect/HLS/Transforms/ScheduleDataflow.cpp b/lib/Dialect/HLS/Transforms/ScheduleDataflow.cpp index 9fae6df1..1bb6a76e 100644 --- a/lib/Dialect/HLS/Transforms/ScheduleDataflow.cpp +++ b/lib/Dialect/HLS/Transforms/ScheduleDataflow.cpp @@ -132,6 +132,19 @@ struct ScheduleDataflow } } + void applyDefaultInstanceLocations() { + SmallVector instances; + getOperation().walk([&](Operation *op) { + if (isa(op)) + instances.push_back(op); + }); + for (auto instance : instances) { + assert(llvm::hasSingleElement(instance->getUsers()) && + "instance should have a single user"); + instance->moveBefore(*instance->user_begin()); + } + } + // Recursively add the task to the group, and create a new group if the task // is not in the group. void dfsScheduleDefiningOp(Value value, size_t prevLevel) { @@ -142,8 +155,8 @@ struct ScheduleDataflow opToLevelMap.lookup(definingOp) > prevLevel) return; - assert(!isa(definingOp) && - !isa(definingOp) && + assert(!isa(definingOp) && + !isa(definingOp) && "tensor/itensor init op should not be scheduled at all"); if (auto task = dyn_cast(definingOp)) { @@ -177,6 +190,17 @@ struct ScheduleDataflow auto func = getOperation(); OpBuilder builder(&getContext()); + // Check if all itensor/tensor init ops have been converted. + auto checkResult = func.walk([&](Operation *op) { + if (isa(op)) { + op->emitOpError("tensor/itensor init op should have been converted"); + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + if (checkResult.wasInterrupted()) + return signalPassFailure(); + // If the task locations are not set, apply the default task locations. if (!checkTaskLocations()) applyDefaultTaskLocations(); @@ -210,6 +234,8 @@ struct ScheduleDataflow func.getName().str() + "_schedule_" + std::to_string(taskId++); wrapOpsIntoTask(ops, taskName, location, builder); } + + applyDefaultInstanceLocations(); } private: