Orbit's Level-Scale-Aware ILP Bootstrap and Rescale Placement#2979
Orbit's Level-Scale-Aware ILP Bootstrap and Rescale Placement#2979edwjchen wants to merge 1 commit into
Conversation
|
I will review this PR this week, sorry for the delay! |
j2kun
left a comment
j2kun
left a comment
There was a problem hiding this comment.
Generally it looks good. I have some comments on organization, one about the integrality of the state variables, and then I'm still a little bit confused on where the rescale is happening. It seems like this model is not yet handling rescale decisions, but instead is using level_reduce to align mis-aligned operand levels (which is good but not quite the same as a rescale decision). It may be that refactoring the code will make it easier for me to see exactly which constraint is controlling the scale updates.
| // ILP Formulation: | ||
| // - Variables: level[value] for each SSA value, input_level[op] for each op, | ||
| // bootstrap[op] for each operation | ||
| // - Variables: level[value] and scale[value] for each SSA value, |
There was a problem hiding this comment.
| // - Variables: level[value] and scale[value] for each SSA value, | |
| // - Variables: level[value] and scale[value] (in bits) for each SSA value, |
nit: may be worth specifying the units on scale.
| // Operand matching: all operands must be at the same level when consumed. | ||
| // inputLevel = level at which operands are consumed. | ||
| // Operand matching: all operands must be consumed at the same level. | ||
| // Unlike the original model, a producer may stay at a higher level and a |
There was a problem hiding this comment.
nit: mentioning "the original model" is not particularly informative, given there is no explanation of why the difference to the original model is useful. (Is "the original model" the earlier implementation? an external impl? best to just not mention it and leave discussion of the changes to the PR/commit description)
| if (!root) return failure(); | ||
|
|
||
| std::optional<int> parsedBootstrapCost = | ||
| averagePositiveLatency(*root, "earth.bootstrap_single"); |
There was a problem hiding this comment.
Is there a hope to make this more fine-grained than using a blended cost? If not, perhaps the json schema itself should just provide a blended latency cost for each op. I think supporting a per-level cost would require pretty dramatic changes to the model.
There was a problem hiding this comment.
It would be helpful to have documented somewhere (perhaps in a README.md in the lib/Transforms/ILPBootstrapPlacement or analysis directory) the specification of the cost model json file, and to include some breadcrumbs (maybe in the pass tablegen) linking back to that document, so the meaning of the fields of this json can be understood in the future.
Also some nits on the json spec itself, which would help to either fix or explain in the readme:
- Why use
earth.*_single, which I think came from DaCapo (?) instead of our own op names? - Why are there multiple values for each op in the latency table? Are these latencies per-level, and if so (a) what is the ordering of the levels and (b) why do they have a different number of values (bootstrap has 16, rescale 15). 16 also seems quite low, so I'm wondering what you'd expect a pass to do when they want the cost of an op at level 30.
- What is the unit of latency?
- How do the other fields influence the rest of the cost model? (They seem unused, so maybe remove them and re-introduce them in the future when they are used)
| SmallVector<ILPBootstrapPlacementAnalysis::OperandLevelReduction, 32>* | ||
| operandLevelReductions) { | ||
| genericOp->walk([&](mgmt::BootstrapOp op) { | ||
| op.getResult().replaceAllUsesWith(op.getOperand()); |
There was a problem hiding this comment.
sanity check: should you also remove any existing mgmt.modreduce and mgmt.level_reduce ops?
| genericOp->emitError() << "failed to load Orbit cost model from `" | ||
| << orbitCostModel << "`"; |
There was a problem hiding this comment.
| genericOp->emitError() << "failed to load Orbit cost model from `" | |
| << orbitCostModel << "`"; |
Nit, unnecessary.
| if (!insertAfter) continue; | ||
|
|
||
| b.setInsertionPointAfter(insertAfter); | ||
| auto levelReduceOp = mgmt::LevelReduceOp::create( |
There was a problem hiding this comment.
mgmt.level_reduce drops the last limb, while mgmt.modreduce is what lowers to a CKKS rescale op. Is level_reduce what you want?
| } | ||
| } | ||
|
|
||
| LogicalResult ILPBootstrapPlacementAnalysis::solve() { |
There was a problem hiding this comment.
Since this function is now 500 lines long, I think it would help greatly to split out the component pieces into separate functions. Maybe one function to construct the model and one to solve it, and then in the model construction function you can split up the relevant sub-tasks into more readable pieces (e.g., generating and indexing the variables, generating each type of constraint, applying compression).
| ss << "level" << opName << result.getResultNumber(); | ||
| auto levelVar = | ||
| model.AddContinuousVariable(0, bootstrapWaterline, ss.str()); | ||
| auto levelVar = model.AddIntegerVariable(0, bootstrapWaterline, ss.str()); |
There was a problem hiding this comment.
If the level input constraints are integral, and the level drop variables are integral, then the level variables can be left continuous and they will force the state variables to be integral (as they were before this change). I don't see any newly added constraints in this PR that would contradict this.
2 participants
The ILP formulation is inspired by Orbit
This PR extends the ILP bootstrap placement model to solve for both bootstrap and rescale placement.
Changes: