3 files changed, 17 insertions, 10 deletions
diff --git a/frontends/ast/ast.cc b/frontends/ast/ast.cc
index 4e61b33a9..d35ea4171 100644
--- a/frontends/ast/ast.cc
+++ b/frontends/ast/ast.cc
@@ -46,7 +46,7 @@ namespace AST {
 
 // instanciate global variables (private API)
 namespace AST_INTERNAL {
-	bool flag_dump_ast, flag_dump_ast_diff, flag_dump_vlog, flag_nolatches, flag_nomem2reg;
+	bool flag_dump_ast, flag_dump_ast_diff, flag_dump_vlog, flag_nolatches, flag_nomem2reg, flag_mem2reg;
 	AstNode *current_ast, *current_ast_mod;
 	std::map<std::string, AstNode*> current_scope;
 	RTLIL::SigSpec *genRTLIL_subst_from = NULL;
@@ -704,11 +704,12 @@ static AstModule* process_module(AstNode *ast)
 	current_module->ast = ast_before_simplify;
 	current_module->nolatches = flag_nolatches;
 	current_module->nomem2reg = flag_nomem2reg;
+	current_module->mem2reg = flag_mem2reg;
 	return current_module;
 }
 
 // create AstModule instances for all modules in the AST tree and add them to 'design'
-void AST::process(RTLIL::Design *design, AstNode *ast, bool dump_ast, bool dump_ast_diff, bool dump_vlog, bool nolatches, bool nomem2reg)
+void AST::process(RTLIL::Design *design, AstNode *ast, bool dump_ast, bool dump_ast_diff, bool dump_vlog, bool nolatches, bool nomem2reg, bool mem2reg)
 {
 	current_ast = ast;
 	flag_dump_ast = dump_ast;
@@ -716,6 +717,7 @@ void AST::process(RTLIL::Design *design, AstNode *ast, bool dump_ast, bool dump_
 	flag_dump_vlog = dump_vlog;
 	flag_nolatches = nolatches;
 	flag_nomem2reg = nomem2reg;
+	flag_mem2reg = mem2reg;
 
 	assert(current_ast->type == AST_DESIGN);
 	for (auto it = current_ast->children.begin(); it != current_ast->children.end(); it++) {
@@ -744,6 +746,7 @@ RTLIL::IdString AstModule::derive(RTLIL::Design *design, std::map<RTLIL::IdStrin
 	flag_dump_vlog = false;
 	flag_nolatches = nolatches;
 	flag_nomem2reg = nomem2reg;
+	flag_mem2reg = mem2reg;
 	use_internal_line_num();
 
 	std::vector<unsigned char> hash_data;
@@ -817,6 +820,7 @@ void AstModule::update_auto_wires(std::map<RTLIL::IdString, int> auto_sizes)
 	flag_dump_vlog = false;
 	flag_nolatches = nolatches;
 	flag_nomem2reg = nomem2reg;
+	flag_mem2reg = mem2reg;
 	use_internal_line_num();
 
 	for (auto it = auto_sizes.begin(); it != auto_sizes.end(); it++) {
diff --git a/frontends/ast/ast.h b/frontends/ast/ast.h
index 76314f88d..81d29a02f 100644
--- a/frontends/ast/ast.h
+++ b/frontends/ast/ast.h
@@ -164,7 +164,7 @@ namespace AST
 		bool simplify(bool const_fold, bool at_zero, bool in_lvalue, int stage);
 		void expand_genblock(std::string index_var, std::string prefix, std::map<std::string, std::string> &name_map);
 		void replace_ids(std::map<std::string, std::string> &rules);
-		void mem2reg_as_needed_pass1(std::set<AstNode*> &mem2reg_set, std::set<AstNode*> &mem2reg_candidates, bool sync_proc, bool async_proc);
+		void mem2reg_as_needed_pass1(std::set<AstNode*> &mem2reg_set, std::set<AstNode*> &mem2reg_candidates, bool sync_proc, bool async_proc, bool force_mem2reg);
 		void mem2reg_as_needed_pass2(std::set<AstNode*> &mem2reg_set, AstNode *mod, AstNode *block, AstNode *top_block);
 		void meminfo(int &mem_width, int &mem_size, int &addr_bits);
 
@@ -189,13 +189,13 @@ namespace AST
 	};
 
 	// process an AST tree (ast must point to an AST_DESIGN node) and generate RTLIL code
-	void process(RTLIL::Design *design, AstNode *ast, bool dump_ast = false, bool dump_ast_diff = false, bool dump_vlog = false, bool nolatches = false, bool nomem2reg = false);
+	void process(RTLIL::Design *design, AstNode *ast, bool dump_ast = false, bool dump_ast_diff = false, bool dump_vlog = false, bool nolatches = false, bool nomem2reg = false, bool mem2reg = false);
 
 	// parametric modules are supported directly by the AST library
 	// therfore we need our own derivate of RTLIL::Module with overloaded virtual functions
 	struct AstModule : RTLIL::Module {
 		AstNode *ast;
-		bool nolatches, nomem2reg;
+		bool nolatches, nomem2reg, mem2reg;
 		virtual ~AstModule();
 		virtual RTLIL::IdString derive(RTLIL::Design *design, std::map<RTLIL::IdString, RTLIL::Const> parameters);
 		virtual void update_auto_wires(std::map<RTLIL::IdString, int> auto_sizes);
@@ -217,7 +217,7 @@ namespace AST
 namespace AST_INTERNAL
 {
 	// internal state variables
-	extern bool flag_dump_ast, flag_dump_ast_diff, flag_nolatches, flag_nomem2reg;
+	extern bool flag_dump_ast, flag_dump_ast_diff, flag_nolatches, flag_nomem2reg, flag_mem2reg;
 	extern AST::AstNode *current_ast, *current_ast_mod;
 	extern std::map<std::string, AST::AstNode*> current_scope;
 	extern RTLIL::SigSpec *genRTLIL_subst_from, *genRTLIL_subst_to;
diff --git a/frontends/ast/simplify.cc b/frontends/ast/simplify.cc
index 6594cfcb3..8a02cc129 100644
--- a/frontends/ast/simplify.cc
+++ b/frontends/ast/simplify.cc
@@ -57,7 +57,7 @@ bool AstNode::simplify(bool const_fold, bool at_zero, bool in_lvalue, int stage)
 		if (!flag_nomem2reg && attributes.count("\\nomem2reg") == 0)
 		{
 			std::set<AstNode*> mem2reg_set, mem2reg_candidates;
-			mem2reg_as_needed_pass1(mem2reg_set, mem2reg_candidates, false, false);
+			mem2reg_as_needed_pass1(mem2reg_set, mem2reg_candidates, false, false, flag_mem2reg);
 
 			for (auto node : mem2reg_set)
 			{
@@ -924,7 +924,7 @@ void AstNode::replace_ids(std::map<std::string, std::string> &rules)
 }
 
 // find memories that should be replaced by registers
-void AstNode::mem2reg_as_needed_pass1(std::set<AstNode*> &mem2reg_set, std::set<AstNode*> &mem2reg_candidates, bool sync_proc, bool async_proc)
+void AstNode::mem2reg_as_needed_pass1(std::set<AstNode*> &mem2reg_set, std::set<AstNode*> &mem2reg_candidates, bool sync_proc, bool async_proc, bool force_mem2reg)
 {
 	if ((type == AST_ASSIGN_LE && async_proc) || (type == AST_ASSIGN_EQ && (sync_proc || async_proc)))
 		if (children[0]->type == AST_IDENTIFIER && children[0]->id2ast && children[0]->id2ast->type == AST_MEMORY &&
@@ -938,9 +938,12 @@ void AstNode::mem2reg_as_needed_pass1(std::set<AstNode*> &mem2reg_set, std::set<
 			mem2reg_candidates.insert(children[0]->id2ast);
 		}
 	
-	if (type == AST_MEMORY && attributes.count("\\mem2reg") > 0)
+	if (type == AST_MEMORY && (attributes.count("\\mem2reg") > 0 || force_mem2reg))
 		mem2reg_set.insert(this);
 
+	if (type == AST_MODULE && attributes.count("\\mem2reg") > 0)
+		force_mem2reg = true;
+
 	if (type == AST_ALWAYS) {
 		for (auto child : children) {
 			if (child->type == AST_POSEDGE || child->type == AST_NEGEDGE)
@@ -950,7 +953,7 @@ void AstNode::mem2reg_as_needed_pass1(std::set<AstNode*> &mem2reg_set, std::set<
 	}
 
 	for (auto child : children)
-		child->mem2reg_as_needed_pass1(mem2reg_set, mem2reg_candidates, sync_proc, async_proc);
+		child->mem2reg_as_needed_pass1(mem2reg_set, mem2reg_candidates, sync_proc, async_proc, force_mem2reg);
 }
 
 // actually replace memories with registers