From 76c562bac5c085e84be7bd48a3a3d044d390c2c3 Mon Sep 17 00:00:00 2001 From: Chris Gyurgyik Date: Thu, 9 Jan 2025 15:44:55 -0800 Subject: [PATCH] [LoopScheduleToCalyx] Lower pipeline register to register writes. (#8048) --- .../circt/Dialect/Calyx/CalyxLoweringUtils.h | 19 +++-- .../LoopScheduleToCalyx.cpp | 70 +++++++++++++------ .../LoopScheduleToCalyx/convert_pipeline.mlir | 44 ++++++++++++ 3 files changed, 106 insertions(+), 27 deletions(-) diff --git a/include/circt/Dialect/Calyx/CalyxLoweringUtils.h b/include/circt/Dialect/Calyx/CalyxLoweringUtils.h index 04cbf32e3f32..8cb05f9c115f 100644 --- a/include/circt/Dialect/Calyx/CalyxLoweringUtils.h +++ b/include/circt/Dialect/Calyx/CalyxLoweringUtils.h @@ -29,6 +29,7 @@ #include "llvm/ADT/TypeSwitch.h" #include "llvm/Support/JSON.h" +#include #include namespace circt { @@ -404,12 +405,13 @@ class ComponentLoweringStateInterface { /// Put the name of the callee and the instance of the call into map. void addInstance(StringRef calleeName, InstanceOp instanceOp); - /// Return the group which evaluates the value v. Optionally, caller may - /// specify the expected type of the group. + /// Returns the evaluating group or None if not found. template - TGroupOp getEvaluatingGroup(Value v) { + std::optional findEvaluatingGroup(Value v) { auto it = valueGroupAssigns.find(v); - assert(it != valueGroupAssigns.end() && "No group evaluating value!"); + if (it == valueGroupAssigns.end()) + return std::nullopt; + if constexpr (std::is_same_v) return it->second; else { @@ -419,6 +421,15 @@ class ComponentLoweringStateInterface { } } + /// Return the group which evaluates the value v. Optionally, caller may + /// specify the expected type of the group. + template + TGroupOp getEvaluatingGroup(Value v) { + std::optional group = findEvaluatingGroup(v); + assert(group.has_value() && "No group evaluating value!"); + return *group; + } + template struct IsFloatingPoint : std::false_type {}; diff --git a/lib/Conversion/LoopScheduleToCalyx/LoopScheduleToCalyx.cpp b/lib/Conversion/LoopScheduleToCalyx/LoopScheduleToCalyx.cpp index c7ea6889bff3..2f51e1a42291 100644 --- a/lib/Conversion/LoopScheduleToCalyx/LoopScheduleToCalyx.cpp +++ b/lib/Conversion/LoopScheduleToCalyx/LoopScheduleToCalyx.cpp @@ -1019,11 +1019,11 @@ class BuildPipelineGroups : public calyx::FuncOpPartialLoweringPattern { // Collect group names for the prologue or epilogue. SmallVector prologueGroups, epilogueGroups; + auto &state = getState(); auto updatePrologueAndEpilogue = [&](calyx::GroupOp group) { // Mark the group for scheduling in the pipeline's block. - getState().addBlockScheduleable(stage->getBlock(), - group); + state.addBlockScheduleable(stage->getBlock(), group); // Add the group to the prologue or epilogue for this stage as // necessary. The goal is to fill the pipeline so it will be in steady @@ -1046,8 +1046,7 @@ class BuildPipelineGroups : public calyx::FuncOpPartialLoweringPattern { // Covers the case where there are no values that need to be passed // through to the next stage, e.g., some intermediary store. for (auto &op : stage.getBodyBlock()) - if (auto group = getState() - .getNonPipelinedGroupFrom(&op)) + if (auto group = state.getNonPipelinedGroupFrom(&op)) updatePrologueAndEpilogue(*group); } @@ -1058,26 +1057,51 @@ class BuildPipelineGroups : public calyx::FuncOpPartialLoweringPattern { // Get the pipeline register for that result. auto pipelineRegister = pipelineRegisters[i]; - // Get the evaluating group for that value. - calyx::GroupInterface evaluatingGroup = - getState().getEvaluatingGroup(value); - - // Remember the final group for this stage result. calyx::GroupOp group; - - // Stitch the register in, depending on whether the group was - // combinational or sequential. - if (auto combGroup = - dyn_cast(evaluatingGroup.getOperation())) - group = - convertCombToSeqGroup(combGroup, pipelineRegister, value, rewriter); - else - group = - replaceGroupRegister(evaluatingGroup, pipelineRegister, rewriter); - - // Replace the stage result uses with the register out. - stage.getResult(i).replaceAllUsesWith(pipelineRegister.getOut()); - + // Get the evaluating group for that value. + std::optional evaluatingGroup = + state.findEvaluatingGroup(value); + if (!evaluatingGroup.has_value()) { + if (auto opStage = + dyn_cast(value.getDefiningOp())) { + // The pipeline register for this input value needs to be discovered. + auto opResult = cast(value); + unsigned int opNumber = opResult.getResultNumber(); + auto &stageRegisters = state.getPipelineRegs(opStage); + calyx::RegisterOp opRegister = stageRegisters.find(opNumber)->second; + value = opRegister.getOut(); // Pass the `out` wire of this register. + } + if (value.getDefiningOp() == nullptr) { + // We add this for any unhandled cases. + llvm::errs() << "unexpected: input value: " << value << ", in stage " + << stage.getStageNumber() << " register " << i + << " is not a register and was not previously " + "evaluated in a Calyx group. Please open an issue.\n"; + return LogicalResult::failure(); + } + // This is a register's `out` value being written to this pipeline + // register. We create a new group to build this assignment. + std::string groupName = state.getUniqueName( + loweringState().blockName(pipelineRegister->getBlock())); + group = calyx::createGroup( + rewriter, state.getComponentOp(), pipelineRegister->getLoc(), + groupName); + calyx::buildAssignmentsForRegisterWrite( + rewriter, group, state.getComponentOp(), pipelineRegister, value); + } else { + // This was previously evaluated. Stitch the register in, depending on + // whether the group was combinational or sequential. + auto combGroup = + dyn_cast(evaluatingGroup->getOperation()); + group = combGroup == nullptr + ? replaceGroupRegister(*evaluatingGroup, pipelineRegister, + rewriter) + : convertCombToSeqGroup(combGroup, pipelineRegister, value, + rewriter); + + // Replace the stage result uses with the register out. + stage.getResult(i).replaceAllUsesWith(pipelineRegister.getOut()); + } updatePrologueAndEpilogue(group); } diff --git a/test/Conversion/LoopScheduleToCalyx/convert_pipeline.mlir b/test/Conversion/LoopScheduleToCalyx/convert_pipeline.mlir index 17cc7390c90a..9bf2e9ad39a8 100644 --- a/test/Conversion/LoopScheduleToCalyx/convert_pipeline.mlir +++ b/test/Conversion/LoopScheduleToCalyx/convert_pipeline.mlir @@ -267,3 +267,47 @@ module { } } +// ----- + +// Pipeline register to pipeline register writes are valid. + +//CHECK: calyx.group @bb0_3 { +//CHECK-NEXT: calyx.assign %stage_1_register_0_reg.in = %stage_0_register_0_reg.out : i32 +//CHECK-NEXT: calyx.assign %stage_1_register_0_reg.write_en = %true : i1 +//CHECK-NEXT: calyx.group_done %stage_1_register_0_reg.done : i1 +//CHECK-NEXT: } +//CHECK-NEXT: calyx.group @bb0_4 { +//CHECK-NEXT: calyx.assign %stage_2_register_0_reg.in = %stage_1_register_0_reg.out : i32 +//CHECK-NEXT: calyx.assign %stage_2_register_0_reg.write_en = %true : i1 +//CHECK-NEXT: calyx.group_done %stage_2_register_0_reg.done : i1 +//CHECK-NEXT: } +module { + func.func @foo(%arg0: i32, %arg1: memref<30x30xi32>) attributes {llvm.linkage = #llvm.linkage} { + %alloca = memref.alloca() : memref<30x30xi32> + cf.br ^bb4 + ^bb4: + %0 = arith.constant 0 : index + %c1_5 = arith.constant 1 : index + loopschedule.pipeline II = 2 trip_count = 20 iter_args(%arg5 = %0) : (index) -> () { + %6 = arith.cmpi ult, %arg5, %0 : index + loopschedule.register %6 : i1 + } do { + %6:2 = loopschedule.pipeline.stage start = 0 { + %10 = memref.load %arg1[%0, %0] : memref<30x30xi32> + %11 = arith.addi %arg5, %c1_5 : index + loopschedule.register %10, %11 : i32, index + } : i32, index + %7 = loopschedule.pipeline.stage start = 1 { + loopschedule.register %6#0 : i32 + } : i32 + %8 = loopschedule.pipeline.stage start = 2 { + loopschedule.register %7 : i32 + } : i32 + loopschedule.terminator iter_args(%6#1), results() : (index) -> () + } + cf.br ^bb6 + ^bb6: + return + } +} +