19 files changed, 38 insertions, 54 deletions
diff --git a/riscv/insns/amoadd_d.h b/riscv/insns/amoadd_d.h
index 9c7c124..6090fbc 100644
--- a/riscv/insns/amoadd_d.h
+++ b/riscv/insns/amoadd_d.h
@@ -1,5 +1,3 @@
 require_extension('A');
 require_rv64;
-reg_t v = MMU.load_uint64(RS1);
-MMU.store_uint64(RS1, RS2 + v);
-WRITE_RD(v);
+WRITE_RD(MMU.amo_uint64(RS1, [&](uint64_t lhs) { return lhs + RS2; }));
diff --git a/riscv/insns/amoadd_w.h b/riscv/insns/amoadd_w.h
index 7ac59b0..2c6471a 100644
--- a/riscv/insns/amoadd_w.h
+++ b/riscv/insns/amoadd_w.h
@@ -1,4 +1,2 @@
 require_extension('A');
-reg_t v = MMU.load_int32(RS1);
-MMU.store_uint32(RS1, RS2 + v);
-WRITE_RD(v);
+WRITE_RD(sext32(MMU.amo_uint32(RS1, [&](uint32_t lhs) { return lhs + RS2; })));
diff --git a/riscv/insns/amoand_d.h b/riscv/insns/amoand_d.h
index 7aa6386..80aea18 100644
--- a/riscv/insns/amoand_d.h
+++ b/riscv/insns/amoand_d.h
@@ -1,5 +1,3 @@
 require_extension('A');
 require_rv64;
-reg_t v = MMU.load_uint64(RS1);
-MMU.store_uint64(RS1, RS2 & v);
-WRITE_RD(v);
+WRITE_RD(MMU.amo_uint64(RS1, [&](uint64_t lhs) { return lhs & RS2; }));
diff --git a/riscv/insns/amoand_w.h b/riscv/insns/amoand_w.h
index 7db2160..f7e1ba7 100644
--- a/riscv/insns/amoand_w.h
+++ b/riscv/insns/amoand_w.h
@@ -1,4 +1,2 @@
 require_extension('A');
-reg_t v = MMU.load_int32(RS1);
-MMU.store_uint32(RS1, RS2 & v);
-WRITE_RD(v);
+WRITE_RD(sext32(MMU.amo_uint32(RS1, [&](uint32_t lhs) { return lhs & RS2; })));
diff --git a/riscv/insns/amomax_d.h b/riscv/insns/amomax_d.h
index 0f6da18..496d8ad 100644
--- a/riscv/insns/amomax_d.h
+++ b/riscv/insns/amomax_d.h
@@ -1,5 +1,3 @@
 require_extension('A');
 require_rv64;
-sreg_t v = MMU.load_int64(RS1);
-MMU.store_uint64(RS1, std::max(sreg_t(RS2),v));
-WRITE_RD(v);
+WRITE_RD(MMU.amo_uint64(RS1, [&](int64_t lhs) { return std::max(lhs, int64_t(RS2)); }));
diff --git a/riscv/insns/amomax_w.h b/riscv/insns/amomax_w.h
index 8c9222b..757bdd2 100644
--- a/riscv/insns/amomax_w.h
+++ b/riscv/insns/amomax_w.h
@@ -1,4 +1,2 @@
 require_extension('A');
-int32_t v = MMU.load_int32(RS1);
-MMU.store_uint32(RS1, std::max(int32_t(RS2),v));
-WRITE_RD(v);
+WRITE_RD(sext32(MMU.amo_uint32(RS1, [&](int32_t lhs) { return std::max(lhs, int32_t(RS2)); })));
diff --git a/riscv/insns/amomaxu_d.h b/riscv/insns/amomaxu_d.h
index 6760f91..12b1733 100644
--- a/riscv/insns/amomaxu_d.h
+++ b/riscv/insns/amomaxu_d.h
@@ -1,5 +1,3 @@
 require_extension('A');
 require_rv64;
-reg_t v = MMU.load_uint64(RS1);
-MMU.store_uint64(RS1, std::max(RS2,v));
-WRITE_RD(v);
+WRITE_RD(MMU.amo_uint64(RS1, [&](uint64_t lhs) { return std::max(lhs, RS2); }));
diff --git a/riscv/insns/amomaxu_w.h b/riscv/insns/amomaxu_w.h
index fc83dc3..538df1c 100644
--- a/riscv/insns/amomaxu_w.h
+++ b/riscv/insns/amomaxu_w.h
@@ -1,4 +1,2 @@
 require_extension('A');
-uint32_t v = MMU.load_int32(RS1);
-MMU.store_uint32(RS1, std::max(uint32_t(RS2),v));
-WRITE_RD((int32_t)v);
+WRITE_RD(sext32(MMU.amo_uint32(RS1, [&](uint32_t lhs) { return std::max(lhs, uint32_t(RS2)); })));
diff --git a/riscv/insns/amomin_d.h b/riscv/insns/amomin_d.h
index 8d08984..725d983 100644
--- a/riscv/insns/amomin_d.h
+++ b/riscv/insns/amomin_d.h
@@ -1,5 +1,3 @@
 require_extension('A');
 require_rv64;
-sreg_t v = MMU.load_int64(RS1);
-MMU.store_uint64(RS1, std::min(sreg_t(RS2),v));
-WRITE_RD(v);
+WRITE_RD(MMU.amo_uint64(RS1, [&](int64_t lhs) { return std::min(lhs, int64_t(RS2)); }));
diff --git a/riscv/insns/amomin_w.h b/riscv/insns/amomin_w.h
index 31a8df8..ee53faa 100644
--- a/riscv/insns/amomin_w.h
+++ b/riscv/insns/amomin_w.h
@@ -1,4 +1,2 @@
 require_extension('A');
-int32_t v = MMU.load_int32(RS1);
-MMU.store_uint32(RS1, std::min(int32_t(RS2),v));
-WRITE_RD(v);
+WRITE_RD(sext32(MMU.amo_uint32(RS1, [&](int32_t lhs) { return std::min(lhs, int32_t(RS2)); })));
diff --git a/riscv/insns/amominu_d.h b/riscv/insns/amominu_d.h
index 8a77edc..15b6c0a 100644
--- a/riscv/insns/amominu_d.h
+++ b/riscv/insns/amominu_d.h
@@ -1,5 +1,3 @@
 require_extension('A');
 require_rv64;
-reg_t v = MMU.load_uint64(RS1);
-MMU.store_uint64(RS1, std::min(RS2,v));
-WRITE_RD(v);
+WRITE_RD(MMU.amo_uint64(RS1, [&](uint64_t lhs) { return std::min(lhs, RS2); }));
diff --git a/riscv/insns/amominu_w.h b/riscv/insns/amominu_w.h
index 2b6aaa3..52e1141 100644
--- a/riscv/insns/amominu_w.h
+++ b/riscv/insns/amominu_w.h
@@ -1,4 +1,2 @@
 require_extension('A');
-uint32_t v = MMU.load_int32(RS1);
-MMU.store_uint32(RS1, std::min(uint32_t(RS2),v));
-WRITE_RD((int32_t)v);
+WRITE_RD(sext32(MMU.amo_uint32(RS1, [&](uint32_t lhs) { return std::min(lhs, uint32_t(RS2)); })));
diff --git a/riscv/insns/amoor_d.h b/riscv/insns/amoor_d.h
index 5a69717..de87627 100644
--- a/riscv/insns/amoor_d.h
+++ b/riscv/insns/amoor_d.h
@@ -1,5 +1,3 @@
 require_extension('A');
 require_rv64;
-reg_t v = MMU.load_uint64(RS1);
-MMU.store_uint64(RS1, RS2 | v);
-WRITE_RD(v);
+WRITE_RD(MMU.amo_uint64(RS1, [&](uint64_t lhs) { return lhs | RS2; }));
diff --git a/riscv/insns/amoor_w.h b/riscv/insns/amoor_w.h
index f5b96b9..3455981 100644
--- a/riscv/insns/amoor_w.h
+++ b/riscv/insns/amoor_w.h
@@ -1,4 +1,2 @@
 require_extension('A');
-reg_t v = MMU.load_int32(RS1);
-MMU.store_uint32(RS1, RS2 | v);
-WRITE_RD(v);
+WRITE_RD(sext32(MMU.amo_uint32(RS1, [&](uint32_t lhs) { return lhs | RS2; })));
diff --git a/riscv/insns/amoswap_d.h b/riscv/insns/amoswap_d.h
index 8cf1411..e1bffde 100644
--- a/riscv/insns/amoswap_d.h
+++ b/riscv/insns/amoswap_d.h
@@ -1,5 +1,3 @@
 require_extension('A');
 require_rv64;
-reg_t v = MMU.load_uint64(RS1);
-MMU.store_uint64(RS1, RS2);
-WRITE_RD(v);
+WRITE_RD(MMU.amo_uint64(RS1, [&](uint64_t lhs) { return RS2; }));
diff --git a/riscv/insns/amoswap_w.h b/riscv/insns/amoswap_w.h
index 0764d59..0f78369 100644
--- a/riscv/insns/amoswap_w.h
+++ b/riscv/insns/amoswap_w.h
@@ -1,4 +1,2 @@
 require_extension('A');
-reg_t v = MMU.load_int32(RS1);
-MMU.store_uint32(RS1, RS2);
-WRITE_RD(v);
+WRITE_RD(sext32(MMU.amo_uint32(RS1, [&](uint32_t lhs) { return RS2; })));
diff --git a/riscv/insns/amoxor_d.h b/riscv/insns/amoxor_d.h
index 3970822..1b3c0bf 100644
--- a/riscv/insns/amoxor_d.h
+++ b/riscv/insns/amoxor_d.h
@@ -1,5 +1,3 @@
 require_extension('A');
 require_rv64;
-reg_t v = MMU.load_uint64(RS1);
-MMU.store_uint64(RS1, RS2 ^ v);
-WRITE_RD(v);
+WRITE_RD(MMU.amo_uint64(RS1, [&](uint64_t lhs) { return lhs ^ RS2; }));
diff --git a/riscv/insns/amoxor_w.h b/riscv/insns/amoxor_w.h
index 9889b64..a1ea82f 100644
--- a/riscv/insns/amoxor_w.h
+++ b/riscv/insns/amoxor_w.h
@@ -1,4 +1,2 @@
 require_extension('A');
-reg_t v = MMU.load_int32(RS1);
-MMU.store_uint32(RS1, RS2 ^ v);
-WRITE_RD(v);
+WRITE_RD(sext32(MMU.amo_uint32(RS1, [&](uint32_t lhs) { return lhs ^ RS2; })));
diff --git a/riscv/mmu.h b/riscv/mmu.h
index 1f8d34b..105908e 100644
--- a/riscv/mmu.h
+++ b/riscv/mmu.h
@@ -105,12 +105,32 @@ public:
         store_slow_path(addr, sizeof(type##_t), (const uint8_t*)&val); \
     }
 
+  // template for functions that perform an atomic memory operation
+  #define amo_func(type) \
+    template<typename op> \
+    type##_t amo_##type(reg_t addr, op f) { \
+      if (addr & (sizeof(type##_t)-1)) \
+        throw trap_store_address_misaligned(addr); \
+      try { \
+        auto lhs = load_##type(addr); \
+        store_##type(addr, f(lhs)); \
+        return lhs; \
+      } catch (trap_load_access_fault& t) { \
+        /* AMO faults should be reported as store faults */ \
+        throw trap_store_access_fault(t.get_badaddr()); \
+      } \
+    }
+
   // store value to memory at aligned address
   store_func(uint8)
   store_func(uint16)
   store_func(uint32)
   store_func(uint64)
 
+  // perform an atomic memory operation at an aligned address
+  amo_func(uint32)
+  amo_func(uint64)
+
   static const reg_t ICACHE_ENTRIES = 1024;
 
   inline size_t icache_index(reg_t addr)