Index: arch/ia32/include/mm/asid.h
===================================================================
--- arch/ia32/include/mm/asid.h	(revision bb68433fea16aec0b0b4240149c9c3e8fb347b34)
+++ arch/ia32/include/mm/asid.h	(revision 7e4e5323698b3059ed08f14198db2c9c1fbc746a)
@@ -38,7 +38,7 @@
 typedef int asid_t;
 
-#define ASID_MAX_ARCH		0
+#define ASID_MAX_ARCH		3
 
-#define asid_install(as)
+#define asid_get()		(ASID_START+1)
 
 #endif
Index: arch/ia64/src/mm/frame.c
===================================================================
--- arch/ia64/src/mm/frame.c	(revision bb68433fea16aec0b0b4240149c9c3e8fb347b34)
+++ arch/ia64/src/mm/frame.c	(revision 7e4e5323698b3059ed08f14198db2c9c1fbc746a)
@@ -44,8 +44,8 @@
 	
 	/*
-         * Workaround to prevent slab allocator from allocating frame 0.
-         * Remove the following statement when the kernel is no longer
+	 * Workaround to prevent slab allocator from allocating frame 0.
+	 * Remove the following statement when the kernel is no longer
 	 * identity mapped.
-         */
+	 */
 	frame_mark_unavailable(0, 1);
 
Index: arch/ppc32/include/mm/asid.h
===================================================================
--- arch/ppc32/include/mm/asid.h	(revision bb68433fea16aec0b0b4240149c9c3e8fb347b34)
+++ arch/ppc32/include/mm/asid.h	(revision 7e4e5323698b3059ed08f14198db2c9c1fbc746a)
@@ -32,7 +32,7 @@
 typedef int asid_t;
 
-#define ASID_MAX_ARCH	0
+#define ASID_MAX_ARCH	3
 
-#define asid_install(as)
+#define asid_get()	(ASID_START+1)
 
 #endif
Index: arch/sparc64/src/mm/frame.c
===================================================================
--- arch/sparc64/src/mm/frame.c	(revision bb68433fea16aec0b0b4240149c9c3e8fb347b34)
+++ arch/sparc64/src/mm/frame.c	(revision 7e4e5323698b3059ed08f14198db2c9c1fbc746a)
@@ -34,5 +34,5 @@
 void frame_arch_init(void)
 {
-	zone_create(0, config.memory_size >> FRAME_WIDTH, ADDR2PFN(ALIGN_UP(config.base + config.kernel_size, FRAME_SIZE)), 0);
+	zone_create(0, config.memory_size >> FRAME_WIDTH, 1, 0);
 
 	/*
Index: genarch/src/mm/asid.c
===================================================================
--- genarch/src/mm/asid.c	(revision bb68433fea16aec0b0b4240149c9c3e8fb347b34)
+++ genarch/src/mm/asid.c	(revision 7e4e5323698b3059ed08f14198db2c9c1fbc746a)
@@ -44,4 +44,8 @@
  * a while.
  *
+ * This code depends on the fact that ASIDS_ALLOCABLE
+ * is greater than number of supported CPUs (i.e. the
+ * amount of concurently active address spaces).
+ *
  * Architectures that don't have hardware support for address
  * spaces do not compile with this file.
@@ -58,7 +62,5 @@
 
 /**
- * asidlock protects both the asids_allocated counter
- * and the list of address spaces that were already
- * assigned ASID.
+ * asidlock protects the asids_allocated counter.
  */
 SPINLOCK_INITIALIZE(asidlock);
@@ -66,17 +68,8 @@
 static count_t asids_allocated = 0;
 
-/**
- * List of address spaces with assigned ASID.
- * When the system runs short of allocable
- * ASIDS, inactive address spaces are guaranteed
- * to be at the beginning of the list.
- */
-LIST_INITIALIZE(as_with_asid_head);
-
-
 /** Allocate free address space identifier.
  *
- * This code depends on the fact that ASIDS_ALLOCABLE
- * is greater than number of supported CPUs.
+ * Interrupts must be disabled and as_lock must be held
+ * prior to this call
  *
  * @return New ASID.
@@ -84,5 +77,4 @@
 asid_t asid_get(void)
 {
-	ipl_t ipl;
 	asid_t asid;
 	link_t *tmp;
@@ -93,5 +85,4 @@
 	 */
 	
-	ipl = interrupts_disable();
 	spinlock_lock(&asidlock);
 	if (asids_allocated == ASIDS_ALLOCABLE) {
@@ -107,9 +98,9 @@
 		 * inactive address space.
 		 */
-		tmp = as_with_asid_head.next;
-		ASSERT(tmp != &as_with_asid_head);
+		ASSERT(!list_empty(&inactive_as_with_asid_head));
+		tmp = inactive_as_with_asid_head.next;
 		list_remove(tmp);
 		
-		as = list_get_instance(tmp, as_t, as_with_asid_link);
+		as = list_get_instance(tmp, as_t, inactive_as_with_asid_link);
 		spinlock_lock(&as->lock);
 
@@ -146,5 +137,4 @@
 	
 	spinlock_unlock(&asidlock);
-	interrupts_restore(ipl);
 	
 	return asid;
@@ -171,56 +161,2 @@
 	interrupts_restore(ipl);
 }
-
-/** Install ASID.
- *
- * This function is to be executed on each address space switch.
- *
- * @param as Address space.
- */
-void asid_install(as_t *as)
-{
-	ipl_t ipl;
-	
-	ipl = interrupts_disable();
-	spinlock_lock(&asidlock);
-	spinlock_lock(&as->lock);
-	
-	if (as->asid != ASID_KERNEL) {
-		if (as->asid != ASID_INVALID) {
-			/*
-			 * This address space has valid ASID.
-			 * Remove 'as' from the list of address spaces
-			 * with assigned ASID, so that it can be later
-			 * appended to the tail of the same list.
-			 * This is to prevent stealing of ASIDs from
-			 * recently installed address spaces.
-			 */
-			list_remove(&as->as_with_asid_link);
-		} else {
-			spinlock_unlock(&as->lock);
-			spinlock_unlock(&asidlock);
-	
-			/*
-			 * This address space doesn't have ASID assigned.
-			 * It was stolen or the address space is being
-			 * installed for the first time.
-			 * Allocate new ASID for it.
-			 */
-			as->asid = asid_get();
-			spinlock_lock(&asidlock);
-			spinlock_lock(&as->lock);
-		}
-	
-		/*
-		 * Now it is sure that 'as' has ASID.
-		 * It is therefore appended to the list
-		 * of address spaces from which it can
-		 * be stolen.
-		 */
-		list_append(&as->as_with_asid_link, &as_with_asid_head);
-	}
-	
-	spinlock_unlock(&as->lock);
-	spinlock_unlock(&asidlock);
-	interrupts_restore(ipl);
-}
Index: genarch/src/mm/page_ht.c
===================================================================
--- genarch/src/mm/page_ht.c	(revision bb68433fea16aec0b0b4240149c9c3e8fb347b34)
+++ genarch/src/mm/page_ht.c	(revision 7e4e5323698b3059ed08f14198db2c9c1fbc746a)
@@ -170,5 +170,5 @@
 
 	if (!hash_table_find(&page_ht, key)) {
-		t = (pte_t *) malloc(sizeof(pte_t));
+		t = (pte_t *) malloc(sizeof(pte_t), FRAME_ATOMIC);
 		ASSERT(t != NULL);
 	
Index: generic/include/mm/as.h
===================================================================
--- generic/include/mm/as.h	(revision bb68433fea16aec0b0b4240149c9c3e8fb347b34)
+++ generic/include/mm/as.h	(revision 7e4e5323698b3059ed08f14198db2c9c1fbc746a)
@@ -77,7 +77,11 @@
 struct as {
 	/** Protected by asidlock. Must be acquired before as->lock. */
-	link_t as_with_asid_link;
+	link_t inactive_as_with_asid_link;
 
 	SPINLOCK_DECLARE(lock);
+
+	/** Number of processors on wich is this address space active. */
+	count_t refcount;
+
 	link_t as_area_head;
 
@@ -97,4 +101,7 @@
 extern as_operations_t *as_operations;
 
+extern spinlock_t as_lock;
+extern link_t inactive_as_with_asid_head;
+
 extern void as_init(void);
 extern as_t *as_create(int flags);
@@ -102,5 +109,5 @@
 extern void as_set_mapping(as_t *as, __address page, __address frame);
 extern int as_page_fault(__address page);
-extern void as_install(as_t *m);
+extern void as_switch(as_t *old, as_t *new);
 
 /* Interface to be implemented by architectures. */
Index: generic/include/mm/asid.h
===================================================================
--- generic/include/mm/asid.h	(revision bb68433fea16aec0b0b4240149c9c3e8fb347b34)
+++ generic/include/mm/asid.h	(revision 7e4e5323698b3059ed08f14198db2c9c1fbc746a)
@@ -46,7 +46,8 @@
 
 extern spinlock_t asidlock;
-extern link_t as_with_asid_head;
 
+#ifndef asid_get
 extern asid_t asid_get(void);
+#endif /* !def asid_get */
 extern void asid_put(asid_t asid);
 
Index: generic/src/mm/as.c
===================================================================
--- generic/src/mm/as.c	(revision bb68433fea16aec0b0b4240149c9c3e8fb347b34)
+++ generic/src/mm/as.c	(revision 7e4e5323698b3059ed08f14198db2c9c1fbc746a)
@@ -57,4 +57,13 @@
 as_operations_t *as_operations = NULL;
 
+/** Address space lock. It protects inactive_as_with_asid_head. */
+SPINLOCK_INITIALIZE(as_lock);
+
+/**
+ * This list contains address spaces that are not active on any
+ * processor and that have valid ASID.
+ */
+LIST_INITIALIZE(inactive_as_with_asid_head);
+
 /** Kernel address space. */
 as_t *AS_KERNEL = NULL;
@@ -80,6 +89,5 @@
 
 	as = (as_t *) malloc(sizeof(as_t), 0);
-
-	list_initialize(&as->as_with_asid_link);
+	link_initialize(&as->inactive_as_with_asid_link);
 	spinlock_initialize(&as->lock, "as_lock");
 	list_initialize(&as->as_area_head);
@@ -90,4 +98,5 @@
 		as->asid = ASID_INVALID;
 	
+	as->refcount = 0;
 	as->page_table = page_table_create(flags);
 
@@ -268,27 +277,71 @@
 }
 
-/** Install address space on CPU.
- *
- * @param as Address space.
- */
-void as_install(as_t *as)
+/** Switch address spaces.
+ *
+ * @param old Old address space or NULL.
+ * @param new New address space.
+ */
+void as_switch(as_t *old, as_t *new)
 {
 	ipl_t ipl;
-	
-	asid_install(as);
+	bool needs_asid = false;
 	
 	ipl = interrupts_disable();
-	spinlock_lock(&as->lock);
-	SET_PTL0_ADDRESS(as->page_table);
-	spinlock_unlock(&as->lock);
+	spinlock_lock(&as_lock);
+
+	/*
+	 * First, take care of the old address space.
+	 */	
+	if (old) {
+		spinlock_lock(&old->lock);
+		ASSERT(old->refcount);
+		if((--old->refcount == 0) && (old != AS_KERNEL)) {
+			/*
+			 * The old address space is no longer active on
+			 * any processor. It can be appended to the
+			 * list of inactive address spaces with assigned
+			 * ASID.
+			 */
+			 ASSERT(old->asid != ASID_INVALID);
+			 list_append(&old->inactive_as_with_asid_link, &inactive_as_with_asid_head);
+		}
+		spinlock_unlock(&old->lock);
+	}
+
+	/*
+	 * Second, prepare the new address space.
+	 */
+	spinlock_lock(&new->lock);
+	if ((new->refcount++ == 0) && (new != AS_KERNEL)) {
+		if (new->asid != ASID_INVALID)
+			list_remove(&new->inactive_as_with_asid_link);
+		else
+			needs_asid = true;	/* defer call to asid_get() until new->lock is released */
+	}
+	SET_PTL0_ADDRESS(new->page_table);
+	spinlock_unlock(&new->lock);
+
+	if (needs_asid) {
+		/*
+		 * Allocation of new ASID was deferred
+		 * until now in order to avoid deadlock.
+		 */
+		asid_t asid;
+		
+		asid = asid_get();
+		spinlock_lock(&new->lock);
+		new->asid = asid;
+		spinlock_unlock(&new->lock);
+	}
+	spinlock_unlock(&as_lock);
 	interrupts_restore(ipl);
-
+	
 	/*
 	 * Perform architecture-specific steps.
 	 * (e.g. write ASID to hardware register etc.)
 	 */
-	as_install_arch(as);
-	
-	AS = as;
+	as_install_arch(new);
+	
+	AS = new;
 }
 
Index: generic/src/proc/scheduler.c
===================================================================
--- generic/src/proc/scheduler.c	(revision bb68433fea16aec0b0b4240149c9c3e8fb347b34)
+++ generic/src/proc/scheduler.c	(revision 7e4e5323698b3059ed08f14198db2c9c1fbc746a)
@@ -328,5 +328,5 @@
 			 * Replace the old one with the new one.
 			 */
-			as_install(as2);
+			as_switch(as1, as2);
 		}
 		TASK = THREAD->task;	
@@ -336,5 +336,5 @@
 
 	#ifdef SCHEDULER_VERBOSE
-	printf("cpu%d: tid %d (priority=%d,ticks=%d,nrdy=%d)\n", CPU->id, THREAD->tid, THREAD->priority, THREAD->ticks, CPU->nrdy);
+	printf("cpu%d: tid %d (priority=%d,ticks=%d,nrdy=%d)\n", CPU->id, THREAD->tid, THREAD->priority, THREAD->ticks, atomic_get(&CPU->nrdy));
 	#endif	
 
