Adjust keep_hierarchy behavior

README.md

@@ -575,6 +575,9 @@ Non-standard or SystemVerilog features for formal verification
   ``@(posedge <netname>)`` or ``@(negedge <netname>)`` when ``<netname>``
   is marked with the ``(* gclk *)`` Verilog attribute.
 
+- The `gate_cost_equivalent` attribute on a module can be used to specify
+  the estimated cost of a module as an equivalent number of basic gate
+  instances.
 
 Supported features from SystemVerilog
 =====================================

passes/hierarchy/keep_hierarchy.cc

@@ -23,20 +23,77 @@
 USING_YOSYS_NAMESPACE
 PRIVATE_NAMESPACE_BEGIN
 
+struct ThresholdHiearchyKeeping {
+	Design *design;
+	CellCosts costs;
+	dict<Module *, int> done;
+	pool<Module *> in_progress;
+	uint64_t threshold;
+
+	ThresholdHiearchyKeeping(Design *design, uint64_t threshold)
+		: design(design), costs(design), threshold(threshold) {}
+
+	uint64_t visit(RTLIL::Module *module) {
+		if (module->has_attribute(ID(gate_cost_equivalent)))
+			return module->attributes[ID(gate_cost_equivalent)].as_int();
+
+		if (module->get_blackbox_attribute())
+			log_error("Missing cost information on instanced blackbox %s\n", log_id(module));
+
+		if (done.count(module))
+			return done.at(module);
+
+		if (in_progress.count(module))
+			log_error("Circular hierarchy\n");
+		in_progress.insert(module);
+
+		uint64_t size = 0;
+		module->has_processes_warn();
+
+		for (auto cell : module->cells()) {
+			if (!cell->type.isPublic()) {
+				size += costs.get(cell);
+			} else {
+				RTLIL::Module *submodule = design->module(cell->type);
+				if (!submodule)
+					log_error("Hierarchy contains unknown module '%s' (instanced as %s in %s)\n",
+							  log_id(cell->type), log_id(cell), log_id(module));
+				size += visit(submodule);
+			}
+		}
+
+		if (size > threshold) {
+			log("Keeping %s (estimated size above threshold: %llu > %llu).\n", log_id(module), size, threshold);
+			module->set_bool_attribute(ID::keep_hierarchy);
+			size = 0;
+		}
+
+		in_progress.erase(module);
+		done[module] = size;
+		return size;
+	}
+};
+
 struct KeepHierarchyPass : public Pass {
-	KeepHierarchyPass() : Pass("keep_hierarchy", "add the keep_hierarchy attribute") {}
+	KeepHierarchyPass() : Pass("keep_hierarchy", "selectively add the keep_hierarchy attribute") {}
 	void help() override
 	{
 		//   |---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|
 		log("\n");
-		log("    keep_hierarchy [options]\n");
+		log("    keep_hierarchy [options] [selection]\n");
 		log("\n");
 		log("Add the keep_hierarchy attribute.\n");
 		log("\n");
 		log("    -min_cost <min_cost>\n");
-		log("        only add the attribute to modules estimated to have more\n");
-		log("        than <min_cost> gates after simple techmapping. Intended\n");
-		log("        for tuning trade-offs between quality and yosys runtime.\n");
+		log("        only add the attribute to modules estimated to have more than <min_cost>\n");
+		log("        gates after simple techmapping. Intended for tuning trade-offs between\n");
+		log("        quality and yosys runtime.\n");
+		log("\n");
+		log("        When evaluating a module's cost, gates which are within a submodule\n");
+		log("        which is marked with the keep_hierarchy attribute are not counted\n");
+		log("        towards the upper module's cost. This applies to both when the attribute\n");
+		log("        was added by this command or was pre-existing.\n");
+		log("\n");
 	}
 	void execute(std::vector<std::string> args, RTLIL::Design *design) override
 	{
@@ -54,16 +111,15 @@ struct KeepHierarchyPass : public Pass {
 		}
 		extra_args(args, argidx, design);
 
-		CellCosts costs(design);
+		if (min_cost) {
+			RTLIL::Module *top = design->top_module();
+			if (!top)
+				log_cmd_error("'-min_cost' mode requires a single top module in the design\n");
 
-		for (auto module : design->selected_modules()) {
-			if (min_cost) {
-				unsigned int cost = costs.get(module);
-				if (cost > min_cost) {
-					log("Marking %s (module too big: %d > %d).\n", log_id(module), cost, min_cost);
-					module->set_bool_attribute(ID::keep_hierarchy);
-				}
-			} else {
+			ThresholdHiearchyKeeping worker(design, min_cost);
+			worker.visit(top);
+		} else {
+			for (auto module : design->selected_modules()) {
 				log("Marking %s.\n", log_id(module));
 				module->set_bool_attribute(ID::keep_hierarchy);
 			}

tests/various/keep_hierarchy.ys

@@ -0,0 +1,53 @@
+read_verilog <<EOF
+(* blackbox *)
+(* gate_cost_equivalent=150 *)
+module macro;
+endmodule
+
+module branch1;
+	macro inst1();
+	macro inst2();
+	macro inst3();
+endmodule
+
+module branch2;
+	macro inst1();
+	macro inst2();
+	macro inst3();
+	macro inst4();
+endmodule
+
+// branch3_submod on its own doesn't meet the threshold
+module branch3_submod();
+	wire [2:0] y;
+	wire [2:0] a;
+	wire [2:0] b;
+	assign y = a * b;
+endmodule
+
+// on the other hand four branch3_submods do
+module branch3;
+	branch3_submod inst1();
+	branch3_submod inst2();
+	branch3_submod inst3();
+	branch3_submod inst4();
+endmodule
+
+// wrapper should have zero cost when branch3 is marked
+// keep_hierarchy
+module branch3_wrapper;
+	branch3 inst();
+endmodule
+
+module top;
+	branch1 inst1();
+	branch2 inst2();
+	branch3_wrapper wrapper();
+endmodule
+EOF
+
+hierarchy -top top
+keep_hierarchy -min_cost 500
+select -assert-mod-count 2 A:keep_hierarchy
+select -assert-any A:keep_hierarchy branch2 %i
+select -assert-any A:keep_hierarchy branch3 %i