MLBEDSW-2378 Set NPU base address in ethosu_init

Change-Id: I1145834000ff81d6e497a8fa77bf997478a80372

4 files changed, 271 insertions, 193 deletions

diff --git a/src/ethosu_common.h b/src/ethosu_common.h
index 7c33470..ce7a663 100644
--- a/src/ethosu_common.h
+++ b/src/ethosu_common.h
@@ -20,6 +20,7 @@
 #define ETHOSU_COMMON_H
 
 #include "ethosu55_interface.h"
+#include "ethosu_device.h"
 
 #if !defined(LOG_ENABLED)
 #define LOG_INFO(format, ...)
@@ -38,11 +39,6 @@
 #define ASSERT(args) assert(args)
 #endif
 
-#if defined(CPU_CORTEX_M55)
-#define NPU_BASE ((uint32_t)0x41700000)
-#else
-#define NPU_BASE ((uint32_t)0x41105000)
-#endif
 #define UNUSED(x) ((void)x)
 
 #define VER_STR(X) VNUM_STR(X)
@@ -57,16 +53,4 @@ static const __attribute__((section("npu_driver_version"))) char driver_version_
 static const __attribute__((section("npu_driver_arch_version"))) char driver_arch_version_str[] =
     VER_STR(NNX_ARCH_VERSION_MAJOR) "." VER_STR(NNX_ARCH_VERSION_MINOR) "." VER_STR(NNX_ARCH_VERSION_PATCH);
 
-static inline uint32_t read_reg(uint32_t address)
-{
-    volatile uint32_t *reg = (uint32_t *)(uintptr_t)(NPU_BASE + address);
-    return *reg;
-}
-
-static inline void write_reg(uint32_t address, uint32_t value)
-{
-    volatile uint32_t *reg = (uint32_t *)(uintptr_t)(NPU_BASE + address);
-    *reg                   = value;
-}
-
 #endif // ETHOSU_COMMON_H
diff --git a/src/ethosu_device.c b/src/ethosu_device.c
index a75a8f6..9bda87d 100644
--- a/src/ethosu_device.c
+++ b/src/ethosu_device.c
@@ -19,6 +19,7 @@
 #include "ethosu_common.h"
 
 #include <assert.h>
+#include <stddef.h>
 #include <stdio.h>
 
 #define BASEP_OFFSET 4
@@ -29,18 +30,26 @@
 static uint32_t stream_length = 0;
 #endif
 
-enum ethosu_error_codes ethosu_dev_init(void)
+enum ethosu_error_codes ethosu_dev_init(struct ethosu_device *dev, const void *base_address)
 {
+#if !defined(ARM_NPU_STUB)
+    dev->base_address = (uintptr_t)base_address;
+#else
+    UNUSED(dev);
+    UNUSED(base_address);
+#endif
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_get_id(struct ethosu_id *id)
+enum ethosu_error_codes ethosu_get_id(struct ethosu_device *dev, struct ethosu_id *id)
 {
     struct id_r _id;
 
 #if !defined(ARM_NPU_STUB)
-    _id.word = read_reg(NPU_REG_ID);
+    _id.word = ethosu_read_reg(dev, NPU_REG_ID);
 #else
+    UNUSED(dev);
+
     _id.word           = 0;
     _id.arch_patch_rev = NNX_ARCH_VERSION_PATCH;
     _id.arch_minor_rev = NNX_ARCH_VERSION_MINOR;
@@ -58,12 +67,14 @@ enum ethosu_error_codes ethosu_get_id(struct ethosu_id *id)
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_get_config(struct ethosu_config *config)
+enum ethosu_error_codes ethosu_get_config(struct ethosu_device *dev, struct ethosu_config *config)
 {
     struct config_r cfg = {.word = 0};
 
 #if !defined(ARM_NPU_STUB)
-    cfg.word = read_reg(NPU_REG_CONFIG);
+    cfg.word = ethosu_read_reg(dev, NPU_REG_CONFIG);
+#else
+    UNUSED(dev);
 #endif
 
     config->macs_per_cc        = cfg.macs_per_cc;
@@ -73,7 +84,8 @@ enum ethosu_error_codes ethosu_get_config(struct ethosu_config *config)
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_run_command_stream(const uint8_t *cmd_stream_ptr,
+enum ethosu_error_codes ethosu_run_command_stream(struct ethosu_device *dev,
+                                                  const uint8_t *cmd_stream_ptr,
                                                   uint32_t cms_length,
                                                   const uint64_t *base_addr,
                                                   int num_base_addr)
@@ -97,16 +109,17 @@ enum ethosu_error_codes ethosu_run_command_stream(const uint8_t *cmd_stream_ptr,
 
     for (int i = 0; i < num_base_reg; i++)
     {
-        write_reg(NPU_REG_BASEP0 + (i * BASEP_OFFSET), reg_basep[i]);
+        ethosu_write_reg(dev, NPU_REG_BASEP0 + (i * BASEP_OFFSET), reg_basep[i]);
     }
 
-    write_reg(NPU_REG_QBASE0, qbase0);
-    write_reg(NPU_REG_QBASE1, qbase1);
-    write_reg(NPU_REG_QSIZE, qsize);
+    ethosu_write_reg(dev, NPU_REG_QBASE0, qbase0);
+    ethosu_write_reg(dev, NPU_REG_QBASE1, qbase1);
+    ethosu_write_reg(dev, NPU_REG_QSIZE, qsize);
 
-    ret_code = ethosu_set_command_run();
+    ret_code = ethosu_set_command_run(dev);
 #else
     // NPU stubbed
+    UNUSED(dev);
     stream_length = cms_length;
     UNUSED(cmd_stream_ptr);
     UNUSED(base_addr);
@@ -119,11 +132,11 @@ enum ethosu_error_codes ethosu_run_command_stream(const uint8_t *cmd_stream_ptr,
     return ret_code;
 }
 
-enum ethosu_error_codes ethosu_is_irq_raised(uint8_t *irq_raised)
+enum ethosu_error_codes ethosu_is_irq_raised(struct ethosu_device *dev, uint8_t *irq_raised)
 {
 #if !defined(ARM_NPU_STUB)
     struct status_r status;
-    status.word = read_reg(NPU_REG_STATUS);
+    status.word = ethosu_read_reg(dev, NPU_REG_STATUS);
     if (status.irq_raised == 1)
     {
         *irq_raised = 1;
@@ -133,30 +146,32 @@ enum ethosu_error_codes ethosu_is_irq_raised(uint8_t *irq_raised)
         *irq_raised = 0;
     }
 #else
+    UNUSED(dev);
     *irq_raised = 1;
 #endif
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_clear_irq_status(void)
+enum ethosu_error_codes ethosu_clear_irq_status(struct ethosu_device *dev)
 {
 #if !defined(ARM_NPU_STUB)
     struct cmd_r oldcmd;
-    oldcmd.word = read_reg(NPU_REG_CMD);
+    oldcmd.word = ethosu_read_reg(dev, NPU_REG_CMD);
 
     struct cmd_r cmd;
     cmd.word           = 0;
     cmd.clear_irq      = 1;
     cmd.clock_q_enable = oldcmd.clock_q_enable;
     cmd.power_q_enable = oldcmd.power_q_enable;
-    write_reg(NPU_REG_CMD, cmd.word);
+    ethosu_write_reg(dev, NPU_REG_CMD, cmd.word);
 #else
+    UNUSED(dev);
 #endif
     return ETHOSU_SUCCESS;
 }
 
 // TODO Understand settings of privilege/security level and update API.
-enum ethosu_error_codes ethosu_soft_reset(void)
+enum ethosu_error_codes ethosu_soft_reset(struct ethosu_device *dev)
 {
     enum ethosu_error_codes return_code = ETHOSU_SUCCESS;
 #if !defined(ARM_NPU_STUB)
@@ -164,10 +179,10 @@ enum ethosu_error_codes ethosu_soft_reset(void)
     struct prot_r prot;
 
     reset.word        = 0;
-    reset.pending_CPL = PRIVILEGE_LEVEL_USER;      // TODO, how to get the host priviledge level
+    reset.pending_CPL = PRIVILEGE_LEVEL_USER;      // TODO, how to get the host privilege level
     reset.pending_CSL = SECURITY_LEVEL_NON_SECURE; // TODO, how to get Security level
 
-    prot.word = read_reg(NPU_REG_PROT);
+    prot.word = ethosu_read_reg(dev, NPU_REG_PROT);
 
     if (prot.active_CPL < reset.pending_CPL && prot.active_CSL > reset.pending_CSL)
     {
@@ -175,15 +190,17 @@ enum ethosu_error_codes ethosu_soft_reset(void)
         return ETHOSU_GENERIC_FAILURE;
     }
     // Reset and set security level
-    write_reg(NPU_REG_RESET, reset.word);
+    ethosu_write_reg(dev, NPU_REG_RESET, reset.word);
 
-    return_code = ethosu_wait_for_reset();
+    return_code = ethosu_wait_for_reset(dev);
+#else
+    UNUSED(dev);
 #endif
 
     return return_code;
 }
 
-enum ethosu_error_codes ethosu_wait_for_reset(void)
+enum ethosu_error_codes ethosu_wait_for_reset(struct ethosu_device *dev)
 {
 #if !defined(ARM_NPU_STUB)
     struct status_r status;
@@ -191,7 +208,7 @@ enum ethosu_error_codes ethosu_wait_for_reset(void)
     // Wait until reset status indicates that reset has been completed
     for (int i = 0; i < 100000; i++)
     {
-        status.word = read_reg(NPU_REG_STATUS);
+        status.word = ethosu_read_reg(dev, NPU_REG_STATUS);
         if (0 == status.reset_status)
         {
             break;
@@ -202,12 +219,17 @@ enum ethosu_error_codes ethosu_wait_for_reset(void)
     {
         return ETHOSU_GENERIC_FAILURE;
     }
+#else
+    UNUSED(dev);
 #endif
 
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_read_apb_reg(uint32_t start_address, uint16_t num_reg, uint32_t *reg)
+enum ethosu_error_codes ethosu_read_apb_reg(struct ethosu_device *dev,
+                                            uint32_t start_address,
+                                            uint16_t num_reg,
+                                            uint32_t *reg)
 {
 #if !defined(ARM_NPU_STUB)
     uint32_t address = start_address;
@@ -216,11 +238,12 @@ enum ethosu_error_codes ethosu_read_apb_reg(uint32_t start_address, uint16_t num
 
     for (int i = 0; i < num_reg; i++)
     {
-        reg[i] = read_reg(address);
+        reg[i] = ethosu_read_reg(dev, address);
         address += REG_OFFSET;
     }
 #else
     // NPU stubbed
+    UNUSED(dev);
     UNUSED(start_address);
     UNUSED(num_reg);
     UNUSED(reg);
@@ -229,22 +252,25 @@ enum ethosu_error_codes ethosu_read_apb_reg(uint32_t start_address, uint16_t num
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_set_qconfig(enum ethosu_memory_type memory_type)
+enum ethosu_error_codes ethosu_set_qconfig(struct ethosu_device *dev, enum ethosu_memory_type memory_type)
 {
     if (memory_type > ETHOSU_AXI1_OUTSTANDING_COUNTER3)
     {
         return ETHOSU_INVALID_PARAM;
     }
 #if !defined(ARM_NPU_STUB)
-    write_reg(NPU_REG_QCONFIG, memory_type);
+    ethosu_write_reg(dev, NPU_REG_QCONFIG, memory_type);
 #else
     // NPU stubbed
+    UNUSED(dev);
     UNUSED(memory_type);
 #endif
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_set_regioncfg(uint8_t region, enum ethosu_memory_type memory_type)
+enum ethosu_error_codes ethosu_set_regioncfg(struct ethosu_device *dev,
+                                             uint8_t region,
+                                             enum ethosu_memory_type memory_type)
 {
     if (region > 7)
     {
@@ -252,19 +278,21 @@ enum ethosu_error_codes ethosu_set_regioncfg(uint8_t region, enum ethosu_memory_
     }
 #if !defined(ARM_NPU_STUB)
     struct regioncfg_r regioncfg;
-    regioncfg.word = read_reg(NPU_REG_REGIONCFG);
+    regioncfg.word = ethosu_read_reg(dev, NPU_REG_REGIONCFG);
     regioncfg.word &= ~(0x3 << (2 * region));
     regioncfg.word |= (memory_type & 0x3) << (2 * region);
-    write_reg(NPU_REG_REGIONCFG, regioncfg.word);
+    ethosu_write_reg(dev, NPU_REG_REGIONCFG, regioncfg.word);
 #else
     // NPU stubbed
+    UNUSED(dev);
     UNUSED(region);
     UNUSED(memory_type);
 #endif
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_set_axi_limit0(enum ethosu_axi_limit_beats max_beats,
+enum ethosu_error_codes ethosu_set_axi_limit0(struct ethosu_device *dev,
+                                              enum ethosu_axi_limit_beats max_beats,
                                               enum ethosu_axi_limit_mem_type memtype,
                                               uint8_t max_reads,
                                               uint8_t max_writes)
@@ -276,9 +304,10 @@ enum ethosu_error_codes ethosu_set_axi_limit0(enum ethosu_axi_limit_beats max_be
     axi_limit0.max_outstanding_read_m1  = max_reads - 1;
     axi_limit0.max_outstanding_write_m1 = max_writes - 1;
 
-    write_reg(NPU_REG_AXI_LIMIT0, axi_limit0.word);
+    ethosu_write_reg(dev, NPU_REG_AXI_LIMIT0, axi_limit0.word);
 #else
     // NPU stubbed
+    UNUSED(dev);
     UNUSED(max_beats);
     UNUSED(memtype);
     UNUSED(max_reads);
@@ -288,7 +317,8 @@ enum ethosu_error_codes ethosu_set_axi_limit0(enum ethosu_axi_limit_beats max_be
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_set_axi_limit1(enum ethosu_axi_limit_beats max_beats,
+enum ethosu_error_codes ethosu_set_axi_limit1(struct ethosu_device *dev,
+                                              enum ethosu_axi_limit_beats max_beats,
                                               enum ethosu_axi_limit_mem_type memtype,
                                               uint8_t max_reads,
                                               uint8_t max_writes)
@@ -300,9 +330,10 @@ enum ethosu_error_codes ethosu_set_axi_limit1(enum ethosu_axi_limit_beats max_be
     axi_limit1.max_outstanding_read_m1  = max_reads - 1;
     axi_limit1.max_outstanding_write_m1 = max_writes - 1;
 
-    write_reg(NPU_REG_AXI_LIMIT1, axi_limit1.word);
+    ethosu_write_reg(dev, NPU_REG_AXI_LIMIT1, axi_limit1.word);
 #else
     // NPU stubbed
+    UNUSED(dev);
     UNUSED(max_beats);
     UNUSED(memtype);
     UNUSED(max_reads);
@@ -312,7 +343,8 @@ enum ethosu_error_codes ethosu_set_axi_limit1(enum ethosu_axi_limit_beats max_be
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_set_axi_limit2(enum ethosu_axi_limit_beats max_beats,
+enum ethosu_error_codes ethosu_set_axi_limit2(struct ethosu_device *dev,
+                                              enum ethosu_axi_limit_beats max_beats,
                                               enum ethosu_axi_limit_mem_type memtype,
                                               uint8_t max_reads,
                                               uint8_t max_writes)
@@ -324,9 +356,10 @@ enum ethosu_error_codes ethosu_set_axi_limit2(enum ethosu_axi_limit_beats max_be
     axi_limit2.max_outstanding_read_m1  = max_reads - 1;
     axi_limit2.max_outstanding_write_m1 = max_writes - 1;
 
-    write_reg(NPU_REG_AXI_LIMIT2, axi_limit2.word);
+    ethosu_write_reg(dev, NPU_REG_AXI_LIMIT2, axi_limit2.word);
 #else
     // NPU stubbed
+    UNUSED(dev);
     UNUSED(max_beats);
     UNUSED(memtype);
     UNUSED(max_reads);
@@ -336,7 +369,8 @@ enum ethosu_error_codes ethosu_set_axi_limit2(enum ethosu_axi_limit_beats max_be
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_set_axi_limit3(enum ethosu_axi_limit_beats max_beats,
+enum ethosu_error_codes ethosu_set_axi_limit3(struct ethosu_device *dev,
+                                              enum ethosu_axi_limit_beats max_beats,
                                               enum ethosu_axi_limit_mem_type memtype,
                                               uint8_t max_reads,
                                               uint8_t max_writes)
@@ -348,9 +382,10 @@ enum ethosu_error_codes ethosu_set_axi_limit3(enum ethosu_axi_limit_beats max_be
     axi_limit3.max_outstanding_read_m1  = max_reads - 1;
     axi_limit3.max_outstanding_write_m1 = max_writes - 1;
 
-    write_reg(NPU_REG_AXI_LIMIT3, axi_limit3.word);
+    ethosu_write_reg(dev, NPU_REG_AXI_LIMIT3, axi_limit3.word);
 #else
     // NPU stubbed
+    UNUSED(dev);
     UNUSED(max_beats);
     UNUSED(memtype);
     UNUSED(max_reads);
@@ -360,102 +395,109 @@ enum ethosu_error_codes ethosu_set_axi_limit3(enum ethosu_axi_limit_beats max_be
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_get_revision(uint32_t *revision)
+enum ethosu_error_codes ethosu_get_revision(struct ethosu_device *dev, uint32_t *revision)
 {
 #if !defined(ARM_NPU_STUB)
-    *revision = read_reg(NPU_REG_REVISION);
+    *revision = ethosu_read_reg(dev, NPU_REG_REVISION);
 #else
-    *revision         = 0xDEADC0DE;
+    UNUSED(dev);
+    *revision = 0xDEADC0DE;
 #endif
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_get_qread(uint32_t *qread)
+enum ethosu_error_codes ethosu_get_qread(struct ethosu_device *dev, uint32_t *qread)
 {
 #if !defined(ARM_NPU_STUB)
-    *qread = read_reg(NPU_REG_QREAD);
+    *qread = ethosu_read_reg(dev, NPU_REG_QREAD);
 #else
-    *qread            = stream_length;
+    UNUSED(dev);
+    *qread = stream_length;
 #endif
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_get_status_mask(uint16_t *status_mask)
+enum ethosu_error_codes ethosu_get_status_mask(struct ethosu_device *dev, uint16_t *status_mask)
 {
 #if !defined(ARM_NPU_STUB)
     struct status_r status;
 
-    status.word  = read_reg(NPU_REG_STATUS);
+    status.word  = ethosu_read_reg(dev, NPU_REG_STATUS);
     *status_mask = status.word & 0xFFFF;
 #else
-    *status_mask      = 0x0000;
+    UNUSED(dev);
+    *status_mask = 0x0000;
 #endif
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_get_irq_history_mask(uint16_t *irq_history_mask)
+enum ethosu_error_codes ethosu_get_irq_history_mask(struct ethosu_device *dev, uint16_t *irq_history_mask)
 {
 #if !defined(ARM_NPU_STUB)
     struct status_r status;
 
-    status.word       = read_reg(NPU_REG_STATUS);
+    status.word       = ethosu_read_reg(dev, NPU_REG_STATUS);
     *irq_history_mask = status.irq_history_mask;
 #else
+    UNUSED(dev);
     *irq_history_mask = 0xffff;
 #endif
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_clear_irq_history_mask(uint16_t irq_history_clear_mask)
+enum ethosu_error_codes ethosu_clear_irq_history_mask(struct ethosu_device *dev, uint16_t irq_history_clear_mask)
 {
 #if !defined(ARM_NPU_STUB)
     struct cmd_r oldcmd;
-    oldcmd.word = read_reg(NPU_REG_CMD);
+    oldcmd.word = ethosu_read_reg(dev, NPU_REG_CMD);
 
     struct cmd_r cmd;
     cmd.word              = 0;
     cmd.clock_q_enable    = oldcmd.clock_q_enable;
     cmd.power_q_enable    = oldcmd.power_q_enable;
     cmd.clear_irq_history = irq_history_clear_mask;
-    write_reg(NPU_REG_CMD, cmd.word);
+    ethosu_write_reg(dev, NPU_REG_CMD, cmd.word);
 #else
+    UNUSED(dev);
     UNUSED(irq_history_clear_mask);
 #endif
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_set_command_run(void)
+enum ethosu_error_codes ethosu_set_command_run(struct ethosu_device *dev)
 {
 #if !defined(ARM_NPU_STUB)
     struct cmd_r oldcmd;
-    oldcmd.word = read_reg(NPU_REG_CMD);
+    oldcmd.word = ethosu_read_reg(dev, NPU_REG_CMD);
 
     struct cmd_r cmd;
     cmd.word                        = 0;
     cmd.transition_to_running_state = 1;
     cmd.clock_q_enable              = oldcmd.clock_q_enable;
     cmd.power_q_enable              = oldcmd.power_q_enable;
-    write_reg(NPU_REG_CMD, cmd.word);
+    ethosu_write_reg(dev, NPU_REG_CMD, cmd.word);
 #else
+    UNUSED(dev);
 #endif
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_get_shram_data(int section, uint32_t *shram_p)
+enum ethosu_error_codes ethosu_get_shram_data(struct ethosu_device *dev, int section, uint32_t *shram_p)
 {
 #if !defined(ARM_NPU_STUB)
     int i            = 0;
     uint32_t address = NPU_REG_SHARED_BUFFER0;
-    write_reg(NPU_REG_DEBUG_ADDRESS, section * BYTES_1KB);
+    ethosu_write_reg(dev, NPU_REG_DEBUG_ADDRESS, section * BYTES_1KB);
 
     while (address <= NPU_REG_SHARED_BUFFER255)
     {
-        shram_p[i] = read_reg(address);
+        shram_p[i] = ethosu_read_reg(dev, address);
         address += REG_OFFSET;
         i++;
     }
 #else
     // NPU stubbed
+    UNUSED(dev);
     UNUSED(section);
     UNUSED(shram_p);
 #endif
@@ -463,7 +505,8 @@ enum ethosu_error_codes ethosu_get_shram_data(int section, uint32_t *shram_p)
     return ETHOSU_SUCCESS;
 }
 
-enum ethosu_error_codes ethosu_set_clock_and_power(enum ethosu_clock_q_request clock_q,
+enum ethosu_error_codes ethosu_set_clock_and_power(struct ethosu_device *dev,
+                                                   enum ethosu_clock_q_request clock_q,
                                                    enum ethosu_power_q_request power_q)
 {
 #if !defined(ARM_NPU_STUB)
@@ -471,10 +514,40 @@ enum ethosu_error_codes ethosu_set_clock_and_power(enum ethosu_clock_q_request c
     cmd.word           = 0;
     cmd.clock_q_enable = clock_q;
     cmd.power_q_enable = power_q;
-    write_reg(NPU_REG_CMD, cmd.word);
+    ethosu_write_reg(dev, NPU_REG_CMD, cmd.word);
 #else
+    UNUSED(dev);
     UNUSED(clock_q);
     UNUSED(power_q);
 #endif
     return ETHOSU_SUCCESS;
 }
+
+uint32_t ethosu_read_reg(struct ethosu_device *dev, uint32_t address)
+{
+#if !defined(ARM_NPU_STUB)
+    ASSERT(dev->base_address != NULL);
+
+    volatile uint32_t *reg = (uint32_t *)(uintptr_t)(dev->base_address + address);
+    return *reg;
+#else
+    UNUSED(dev);
+    UNUSED(address);
+
+    return 0;
+#endif
+}
+
+void ethosu_write_reg(struct ethosu_device *dev, uint32_t address, uint32_t value)
+{
+#if !defined(ARM_NPU_STUB)
+    ASSERT(dev->base_address != NULL);
+
+    volatile uint32_t *reg = (uint32_t *)(uintptr_t)(dev->base_address + address);
+    *reg                   = value;
+#else
+    UNUSED(dev);
+    UNUSED(address);
+    UNUSED(value);
+#endif
+}
diff --git a/src/ethosu_driver.c b/src/ethosu_driver.c
index e372312..1a9337e 100644
--- a/src/ethosu_driver.c
+++ b/src/ethosu_driver.c
@@ -17,9 +17,8 @@
  */
 
 #include "ethosu_driver.h"
-#include "ethosu_config.h"
-
 #include "ethosu_common.h"
+#include "ethosu_config.h"
 #include "ethosu_device.h"
 
 #include <assert.h>
@@ -30,8 +29,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 
-// Abort flag
-static int abort_inference = false;
+struct ethosu_driver ethosu_drv = {.dev = {.base_address = NULL}, .abort_inference = false};
 
 // IRQ
 static volatile bool irq_triggered = false;
@@ -40,20 +38,20 @@ static volatile bool irq_triggered = false;
 void ethosu_irq_handler(void)
 {
     uint8_t irq_raised = 0;
-    (void)ethosu_is_irq_raised(&irq_raised);
+    (void)ethosu_is_irq_raised(&ethosu_drv.dev, &irq_raised);
     ASSERT(irq_raised == 1);
     irq_triggered = true;
-    (void)ethosu_clear_irq_status();
-    (void)ethosu_is_irq_raised(&irq_raised);
+    (void)ethosu_clear_irq_status(&ethosu_drv.dev);
+    (void)ethosu_is_irq_raised(&ethosu_drv.dev, &irq_raised);
     ASSERT(irq_raised == 0);
 }
 
-static inline void wait_for_irq(void)
+static inline void wait_for_irq(struct ethosu_driver *drv)
 {
     while (1)
     {
         __disable_irq();
-        if (irq_triggered || abort_inference)
+        if (irq_triggered || drv->abort_inference)
         {
             __enable_irq();
             break;
@@ -66,13 +64,13 @@ static inline void wait_for_irq(void)
 }
 #else
 // Just polling the status register
-static inline void wait_for_irq(void)
+static inline void wait_for_irq(struct ethosu_driver *drv)
 {
     uint8_t irq_raised = 0;
 
     for (int i = 0; i < 5000; ++i)
     {
-        (void)ethosu_is_irq_raised(&irq_raised);
+        (void)ethosu_is_irq_raised(&drv->dev, &irq_raised);
         if (1 == irq_raised)
         {
             break;
@@ -177,37 +175,44 @@ struct opt_cfg_s
     };
 };
 
-static int handle_optimizer_config(struct opt_cfg_s *opt_cfg_p);
-static int handle_command_stream(const uint8_t *cmd_stream,
+static int handle_optimizer_config(struct ethosu_driver *drv, struct opt_cfg_s *opt_cfg_p);
+static int handle_command_stream(struct ethosu_driver *drv,
+                                 const uint8_t *cmd_stream,
                                  const int cms_length,
                                  const uint64_t *base_addr,
                                  const int num_base_addr);
-static int read_apb_reg(uint16_t);
-static int dump_shram();
-static void dump_npu_register(int npu_reg, int npu_reg_end);
+static int read_apb_reg(struct ethosu_driver *drv, uint16_t);
+static int dump_shram(struct ethosu_driver *drv);
+static void dump_npu_register(struct ethosu_driver *drv, int npu_reg, int npu_reg_end);
 static void dump_command_stream(const uint32_t *cmd_stream, const int cms_length, int qread);
+static void npu_axi_init(struct ethosu_driver *drv);
 
-int ethosu_init(void)
+int ethosu_init(const void *base_address)
 {
     int return_code = 0;
+
     LOG_INFO("ethosu_init calling NPU embed driver ethosu_dev_init\n");
 
-    if (ETHOSU_SUCCESS != ethosu_set_clock_and_power(ETHOSU_CLOCK_Q_DISABLE, ETHOSU_POWER_Q_DISABLE))
+    if (ETHOSU_SUCCESS != ethosu_dev_init(&ethosu_drv.dev, base_address))
+    {
+        LOG_ERR("Failed in ethosu_dev_init");
+        return -1;
+    }
+
+    if (ETHOSU_SUCCESS != ethosu_set_clock_and_power(&ethosu_drv.dev, ETHOSU_CLOCK_Q_DISABLE, ETHOSU_POWER_Q_DISABLE))
     {
         LOG_ERR("Failed to disable clock-q & power-q for Ethos-U\n");
         return -1;
     }
 
-    ethosu_soft_reset();
+    ethosu_soft_reset(&ethosu_drv.dev);
 
-    if (ETHOSU_SUCCESS != ethosu_wait_for_reset())
+    if (ETHOSU_SUCCESS != ethosu_wait_for_reset(&ethosu_drv.dev))
     {
         LOG_ERR("Failed reset of Ethos-U\n");
         return -1;
     }
 
-    return_code = ethosu_dev_init();
-
     return return_code;
 }
 
@@ -219,8 +224,8 @@ int ethosu_get_version(struct ethosu_version *version)
     {
         struct ethosu_id id;
         struct ethosu_config cfg;
-        (void)ethosu_get_id(&id);
-        (void)ethosu_get_config(&cfg);
+        (void)ethosu_get_id(&ethosu_drv.dev, &id);
+        (void)ethosu_get_config(&ethosu_drv.dev, &cfg);
 
         version->id.version_status      = id.version_status;
         version->id.version_minor       = id.version_minor;
@@ -270,7 +275,7 @@ int ethosu_invoke(const void *custom_data_ptr,
     }
     int custom_data_32bit_size = (custom_data_size / BYTES_IN_32_BITS - CUSTOM_OPTION_LENGTH_32_BIT_WORD);
 
-    ethosu_set_clock_and_power(ETHOSU_CLOCK_Q_ENABLE, ETHOSU_POWER_Q_DISABLE);
+    ethosu_set_clock_and_power(&ethosu_drv.dev, ETHOSU_CLOCK_Q_ENABLE, ETHOSU_POWER_Q_DISABLE);
     while (data_ptr < (data_start_ptr + custom_data_32bit_size))
     {
         int ret = 0;
@@ -280,7 +285,7 @@ int ethosu_invoke(const void *custom_data_ptr,
             LOG_INFO("ethosu_invoke OPTIMIZER_CONFIG\n");
             struct opt_cfg_s *opt_cfg_p = (struct opt_cfg_s *)data_ptr;
 
-            ret = handle_optimizer_config(opt_cfg_p);
+            ret = handle_optimizer_config(&ethosu_drv, opt_cfg_p);
             data_ptr += DRIVER_ACTION_LENGTH_32_BIT_WORD + OPTIMIZER_CONFIG_LENGTH_32_BIT_WORD;
             break;
         case COMMAND_STREAM:
@@ -288,33 +293,33 @@ int ethosu_invoke(const void *custom_data_ptr,
             void *command_stream = (uint8_t *)(data_ptr) + sizeof(struct custom_data_s);
             int cms_length       = (data_ptr->reserved << 16) | data_ptr->length;
 
-            abort_inference = false;
+            ethosu_drv.abort_inference = false;
             // It is safe to clear this flag without atomic, because npu is not running.
             irq_triggered = false;
 
-            ret = handle_command_stream(command_stream, cms_length, base_addr, num_base_addr);
+            ret = handle_command_stream(&ethosu_drv, command_stream, cms_length, base_addr, num_base_addr);
 
-            if (ret == -1 && abort_inference)
+            if (return_code == -1 && ethosu_drv.abort_inference)
             {
                 uint32_t qread = 0;
-                ethosu_get_qread(&qread);
+                ethosu_get_qread(&ethosu_drv.dev, &qread);
                 LOG_ERR("NPU timeout\n");
                 dump_command_stream(command_stream, cms_length, qread);
-                dump_npu_register(0x200, 0x2BF);
-                dump_npu_register(0x800, 0xB3F);
-                dump_shram();
+                dump_npu_register(&ethosu_drv, 0x200, 0x2BF);
+                dump_npu_register(&ethosu_drv, 0x800, 0xB3F);
+                dump_shram(&ethosu_drv);
             }
 
             data_ptr += DRIVER_ACTION_LENGTH_32_BIT_WORD + cms_length;
             break;
         case READ_APB_REG:
             LOG_INFO("ethosu_invoke READ_APB_REG\n");
-            ret = read_apb_reg(data_ptr->driver_action_data);
+            ret = read_apb_reg(&ethosu_drv, data_ptr->driver_action_data);
             data_ptr += DRIVER_ACTION_LENGTH_32_BIT_WORD;
             break;
         case DUMP_SHRAM:
             LOG_INFO("ethosu_invoke DUMP_SHRAM\n");
-            ret = dump_shram();
+            ret = dump_shram(&ethosu_drv);
             data_ptr += DRIVER_ACTION_LENGTH_32_BIT_WORD;
             break;
         case NOP:
@@ -332,16 +337,16 @@ int ethosu_invoke(const void *custom_data_ptr,
             break;
         }
     }
-    ethosu_set_clock_and_power(ETHOSU_CLOCK_Q_ENABLE, ETHOSU_POWER_Q_ENABLE);
+    ethosu_set_clock_and_power(&ethosu_drv.dev, ETHOSU_CLOCK_Q_ENABLE, ETHOSU_POWER_Q_ENABLE);
     return return_code;
 }
 
 void ethosu_abort(void)
 {
-    abort_inference = true;
+    ethosu_drv.abort_inference = true;
 }
 
-static int handle_optimizer_config(struct opt_cfg_s *opt_cfg_p)
+static int handle_optimizer_config(struct ethosu_driver *drv, struct opt_cfg_s *opt_cfg_p)
 {
     struct ethosu_config cfg;
     struct ethosu_id id;
@@ -358,8 +363,8 @@ static int handle_optimizer_config(struct opt_cfg_s *opt_cfg_p)
              opt_cfg_p->arch_minor_rev,
              opt_cfg_p->arch_patch_rev);
 
-    (void)ethosu_get_config(&cfg);
-    (void)ethosu_get_id(&id);
+    (void)ethosu_get_config(&drv->dev, &cfg);
+    (void)ethosu_get_id(&drv->dev, &id);
     LOG_INFO("Ethos-U config cmd_stream_version: %d macs_per_cc: %d shram_size: %d\n",
              cfg.cmd_stream_version,
              cfg.macs_per_cc,
@@ -409,38 +414,43 @@ static int handle_optimizer_config(struct opt_cfg_s *opt_cfg_p)
     return return_code;
 }
 
-void npu_axi_init()
+static void npu_axi_init(struct ethosu_driver *drv)
 {
-    ethosu_set_qconfig(NPU_QCONFIG);
-
-    ethosu_set_regioncfg(0, NPU_REGIONCFG_0);
-    ethosu_set_regioncfg(1, NPU_REGIONCFG_1);
-    ethosu_set_regioncfg(2, NPU_REGIONCFG_2);
-    ethosu_set_regioncfg(3, NPU_REGIONCFG_3);
-    ethosu_set_regioncfg(4, NPU_REGIONCFG_4);
-    ethosu_set_regioncfg(5, NPU_REGIONCFG_5);
-    ethosu_set_regioncfg(6, NPU_REGIONCFG_6);
-    ethosu_set_regioncfg(7, NPU_REGIONCFG_7);
-
-    (void)ethosu_set_axi_limit0(AXI_LIMIT0_MAX_BEATS_BYTES,
+    ethosu_set_qconfig(&drv->dev, NPU_QCONFIG);
+
+    ethosu_set_regioncfg(&drv->dev, 0, NPU_REGIONCFG_0);
+    ethosu_set_regioncfg(&drv->dev, 1, NPU_REGIONCFG_1);
+    ethosu_set_regioncfg(&drv->dev, 2, NPU_REGIONCFG_2);
+    ethosu_set_regioncfg(&drv->dev, 3, NPU_REGIONCFG_3);
+    ethosu_set_regioncfg(&drv->dev, 4, NPU_REGIONCFG_4);
+    ethosu_set_regioncfg(&drv->dev, 5, NPU_REGIONCFG_5);
+    ethosu_set_regioncfg(&drv->dev, 6, NPU_REGIONCFG_6);
+    ethosu_set_regioncfg(&drv->dev, 7, NPU_REGIONCFG_7);
+
+    (void)ethosu_set_axi_limit0(&drv->dev,
+                                AXI_LIMIT0_MAX_BEATS_BYTES,
                                 AXI_LIMIT0_MEM_TYPE,
                                 AXI_LIMIT0_MAX_OUTSTANDING_READS,
                                 AXI_LIMIT0_MAX_OUTSTANDING_WRITES);
-    (void)ethosu_set_axi_limit1(AXI_LIMIT1_MAX_BEATS_BYTES,
+    (void)ethosu_set_axi_limit1(&drv->dev,
+                                AXI_LIMIT1_MAX_BEATS_BYTES,
                                 AXI_LIMIT1_MEM_TYPE,
                                 AXI_LIMIT1_MAX_OUTSTANDING_READS,
                                 AXI_LIMIT1_MAX_OUTSTANDING_WRITES);
-    (void)ethosu_set_axi_limit2(AXI_LIMIT2_MAX_BEATS_BYTES,
+    (void)ethosu_set_axi_limit2(&drv->dev,
+                                AXI_LIMIT2_MAX_BEATS_BYTES,
                                 AXI_LIMIT2_MEM_TYPE,
                                 AXI_LIMIT2_MAX_OUTSTANDING_READS,
                                 AXI_LIMIT2_MAX_OUTSTANDING_WRITES);
-    (void)ethosu_set_axi_limit3(AXI_LIMIT3_MAX_BEATS_BYTES,
+    (void)ethosu_set_axi_limit3(&drv->dev,
+                                AXI_LIMIT3_MAX_BEATS_BYTES,
                                 AXI_LIMIT3_MEM_TYPE,
                                 AXI_LIMIT3_MAX_OUTSTANDING_READS,
                                 AXI_LIMIT3_MAX_OUTSTANDING_WRITES);
 }
 
-static int handle_command_stream(const uint8_t *cmd_stream,
+static int handle_command_stream(struct ethosu_driver *drv,
+                                 const uint8_t *cmd_stream,
                                  const int cms_length,
                                  const uint64_t *base_addr,
                                  const int num_base_addr)
@@ -468,16 +478,16 @@ static int handle_command_stream(const uint8_t *cmd_stream,
     {
         return -1;
     }
-    npu_axi_init();
+    npu_axi_init(drv);
 
-    if (ETHOSU_SUCCESS != ethosu_run_command_stream(cmd_stream, cms_bytes, base_addr, num_base_addr))
+    if (ETHOSU_SUCCESS != ethosu_run_command_stream(&drv->dev, cmd_stream, cms_bytes, base_addr, num_base_addr))
     {
         return -1;
     }
 
-    wait_for_irq();
+    wait_for_irq(drv);
 
-    (void)ethosu_get_qread(&qread);
+    (void)ethosu_get_qread(&drv->dev, &qread);
     if (qread != cms_bytes)
     {
         LOG_ERR("Failure: IRQ received but qread (%d) not at end of stream (%d).\n", qread, cms_bytes);
@@ -488,7 +498,7 @@ static int handle_command_stream(const uint8_t *cmd_stream,
     return 0;
 }
 
-static int read_apb_reg(uint16_t da_data)
+static int read_apb_reg(struct ethosu_driver *drv, uint16_t da_data)
 {
     uint32_t *reg_p;
     uint32_t start_address = (uint32_t)(da_data & APB_START_ADDR_MASK);
@@ -501,7 +511,7 @@ static int read_apb_reg(uint16_t da_data)
         return -1;
     }
 
-    if (ETHOSU_SUCCESS == ethosu_read_apb_reg(start_address, num_reg, reg_p))
+    if (ETHOSU_SUCCESS == ethosu_read_apb_reg(&drv->dev, start_address, num_reg, reg_p))
     {
         for (int i = 0; i < num_reg; i++)
         {
@@ -518,11 +528,11 @@ static int read_apb_reg(uint16_t da_data)
     return 0;
 }
 
-static int dump_shram()
+static int dump_shram(struct ethosu_driver *drv)
 {
     struct ethosu_config cfg;
     uint32_t *shram_p;
-    (void)ethosu_get_config(&cfg);
+    (void)ethosu_get_config(&drv->dev, &cfg);
 
     LOG_INFO("dump_shram size = %d KB\n", cfg.shram_size);
 
@@ -535,7 +545,7 @@ static int dump_shram()
 
     for (uint32_t i = 0; i < cfg.shram_size; i++)
     {
-        ethosu_get_shram_data(i, (uint32_t *)shram_p);
+        ethosu_get_shram_data(&drv->dev, i, (uint32_t *)shram_p);
         // Output 1KB of SHRAM
         LOG_INFO("***SHRAM SECTION %d***\n", i);
         for (int j = 0; j < (BYTES_1KB / BYTES_IN_32_BITS); j++)
@@ -725,7 +735,7 @@ static const char *lookup_name(const name_lookup_t *lookup_table, int lookup_tab
     return 0;
 }
 
-static void dump_npu_register(int npu_reg, int npu_reg_end)
+static void dump_npu_register(struct ethosu_driver *drv, int npu_reg, int npu_reg_end)
 {
     unsigned int reg_val;
     const char *reg_name;
@@ -734,7 +744,7 @@ static void dump_npu_register(int npu_reg, int npu_reg_end)
     LOG_INFO("dump_register %X - %X\n", npu_reg, npu_reg_end);
     for (; npu_reg <= npu_reg_end; npu_reg += sizeof(int))
     {
-        reg_val  = read_reg(npu_reg);
+        reg_val  = ethosu_read_reg(&drv->dev, npu_reg);
         reg_name = lookup_name(npu_reg_name_tbl, npu_reg_name_tbl_count, npu_reg);
         LOG_INFO("[0x%.4X] 0x%.8X\t%s\n", npu_reg, reg_val, (reg_name) ? reg_name : "");
     }
diff --git a/src/ethosu_pmu.c b/src/ethosu_pmu.c
index 1d3b66c..62dc5da 100644
--- a/src/ethosu_pmu.c
+++ b/src/ethosu_pmu.c
@@ -22,9 +22,11 @@
 
 #include "ethosu55_interface.h"
 #include "ethosu_common.h"
+#include "ethosu_driver.h"
+#include "pmu_ethosu.h"
+
 #include <assert.h>
-#include <ethosu_driver.h>
-#include <pmu_ethosu.h>
+#include <stddef.h>
 #include <stdint.h>
 
 /*****************************************************************************
@@ -49,6 +51,13 @@
  * Variables
  *****************************************************************************/
 
+/**
+ *  NOTE: A pointer to ethosu_driver will be added to the PMU functions
+ * when multi-NPU functionality is implemented later. We shall use a
+ * shared ethosu_driver instance till then.
+ * */
+extern struct ethosu_driver ethosu_drv;
+
 static const enum pmu_event_type eventbyid[] = {EXPAND_PMU_EVENT_TYPE(EVID, COMMA)};
 
 /*****************************************************************************
@@ -78,21 +87,21 @@ uint32_t pmu_event_value(enum ethosu_pmu_event_type event)
 void ethosu_pmu_driver_init(void)
 {
 #ifdef PMU_AUTOINIT
-    write_reg(NPU_REG_PMCR, INIT_PMCR);
-    write_reg(NPU_REG_PMCNTENSET, INIT_PMCNTENSET);
-    write_reg(NPU_REG_PMCNTENCLR, INIT_PMCNTENCLR);
-    write_reg(NPU_REG_PMOVSSET, INIT_PMOVSSET);
-    write_reg(NPU_REG_PMOVSCLR, INIT_PMOVSCLR);
-    write_reg(NPU_REG_PMINTSET, INIT_PMINTSET);
-    write_reg(NPU_REG_PMINTCLR, INIT_PMINTCLR);
-    write_reg(NPU_REG_PMCCNTR_LO, INIT_PMCCNTR);
-    write_reg(NPU_REG_PMCCNTR_HI, INIT_PMCCNTR);
-    write_reg(NPU_REG_PMCCNTR_CFG, INIT_PMCCNTR_CFG);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCR, INIT_PMCR);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCNTENSET, INIT_PMCNTENSET);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCNTENCLR, INIT_PMCNTENCLR);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMOVSSET, INIT_PMOVSSET);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMOVSCLR, INIT_PMOVSCLR);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMINTSET, INIT_PMINTSET);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMINTCLR, INIT_PMINTCLR);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_LO, INIT_PMCCNTR);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_HI, INIT_PMCCNTR);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_CFG, INIT_PMCCNTR_CFG);
 
     for (int i = 0; i < ETHOSU_PMU_NCOUNTERS; i++)
     {
-        write_reg(NPU_REG_PMEVCNTR(i), 0);
-        write_reg(NPU_REG_PMEVTYPER(i), 0);
+        ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMEVCNTR(i), 0);
+        ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMEVTYPER(i), 0);
     }
 #endif
 }
@@ -102,64 +111,66 @@ void ethosu_pmu_driver_exit(void) {}
 void ETHOSU_PMU_Enable(void)
 {
     struct pmcr_r pmcr;
-    pmcr.word   = read_reg(NPU_REG_PMCR);
+    pmcr.word   = ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMCR);
     pmcr.cnt_en = 1;
-    write_reg(NPU_REG_PMCR, pmcr.word);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCR, pmcr.word);
 }
 
 void ETHOSU_PMU_Disable(void)
 {
     struct pmcr_r pmcr;
-    pmcr.word   = read_reg(NPU_REG_PMCR);
+    pmcr.word   = ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMCR);
     pmcr.cnt_en = 0;
-    write_reg(NPU_REG_PMCR, pmcr.word);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCR, pmcr.word);
 }
 
 void ETHOSU_PMU_Set_EVTYPER(uint32_t num, enum ethosu_pmu_event_type type)
 {
-    write_reg(NPU_REG_PMEVTYPER(num), pmu_event_value(type));
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMEVTYPER(num), pmu_event_value(type));
 }
 
 enum ethosu_pmu_event_type ETHOSU_PMU_Get_EVTYPER(uint32_t num)
 {
-    return pmu_event_type(read_reg(NPU_REG_PMEVTYPER(num)));
+    return pmu_event_type(ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMEVTYPER(num)));
 }
 
 void ETHOSU_PMU_CYCCNT_Reset(void)
 {
     struct pmcr_r pmcr;
-    pmcr.word          = read_reg(NPU_REG_PMCR);
+    pmcr.word          = ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMCR);
     pmcr.cycle_cnt_rst = 1;
-    write_reg(NPU_REG_PMCR, pmcr.word);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCR, pmcr.word);
 }
 
 void ETHOSU_PMU_EVCNTR_ALL_Reset(void)
 {
     struct pmcr_r pmcr;
-    pmcr.word          = read_reg(NPU_REG_PMCR);
+    pmcr.word          = ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMCR);
     pmcr.event_cnt_rst = 1;
-    write_reg(NPU_REG_PMCR, pmcr.word);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCR, pmcr.word);
 }
 
 void ETHOSU_PMU_CNTR_Enable(uint32_t mask)
 {
-    write_reg(NPU_REG_PMCNTENSET, mask);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCNTENSET, mask);
 }
 
 void ETHOSU_PMU_CNTR_Disable(uint32_t mask)
 {
-    write_reg(NPU_REG_PMCNTENCLR, mask);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCNTENCLR, mask);
 }
 
-uint32_t ETHOSU_PMU_CNTR_Status()
+uint32_t ETHOSU_PMU_CNTR_Status(void)
 {
-    return read_reg(NPU_REG_PMCNTENSET);
+    return ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMCNTENSET);
 }
 
 uint64_t ETHOSU_PMU_Get_CCNTR(void)
 {
-    uint64_t val1 = (((uint64_t)read_reg(NPU_REG_PMCCNTR_HI)) << 32) | read_reg(NPU_REG_PMCCNTR_LO);
-    uint64_t val2 = (((uint64_t)read_reg(NPU_REG_PMCCNTR_HI)) << 32) | read_reg(NPU_REG_PMCCNTR_LO);
+    uint64_t val1 = (((uint64_t)ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_HI)) << 32) |
+                    ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_LO);
+    uint64_t val2 = (((uint64_t)ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_HI)) << 32) |
+                    ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_LO);
 
     if (val2 > val1)
     {
@@ -177,8 +188,8 @@ void ETHOSU_PMU_Set_CCNTR(uint64_t val)
         ETHOSU_PMU_CNTR_Disable(ETHOSU_PMU_CCNT_Msk);
     }
 
-    write_reg(NPU_REG_PMCCNTR_LO, (val & MASK_0_31_BITS));
-    write_reg(NPU_REG_PMCCNTR_HI, (val & MASK_32_47_BITS) >> 32);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_LO, (val & MASK_0_31_BITS));
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_HI, (val & MASK_32_47_BITS) >> 32);
 
     if (mask & ETHOSU_PMU_CCNT_Msk)
     {
@@ -188,39 +199,39 @@ void ETHOSU_PMU_Set_CCNTR(uint64_t val)
 
 uint32_t ETHOSU_PMU_Get_EVCNTR(uint32_t num)
 {
-    return read_reg(NPU_REG_PMEVCNTR(num));
+    return ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMEVCNTR(num));
 }
 
 void ETHOSU_PMU_Set_EVCNTR(uint32_t num, uint32_t val)
 {
-    write_reg(NPU_REG_PMEVCNTR(num), val);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMEVCNTR(num), val);
 }
 
 uint32_t ETHOSU_PMU_Get_CNTR_OVS(void)
 {
-    return read_reg(NPU_REG_PMOVSSET);
+    return ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMOVSSET);
 }
 
 // TODO: check if this function name match with the description &
 // implementation.
 void ETHOSU_PMU_Set_CNTR_OVS(uint32_t mask)
 {
-    write_reg(NPU_REG_PMOVSCLR, mask);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMOVSCLR, mask);
 }
 
 void ETHOSU_PMU_Set_CNTR_IRQ_Enable(uint32_t mask)
 {
-    write_reg(NPU_REG_PMINTSET, mask);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMINTSET, mask);
 }
 
 void ETHOSU_PMU_Set_CNTR_IRQ_Disable(uint32_t mask)
 {
-    write_reg(NPU_REG_PMINTCLR, mask);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMINTCLR, mask);
 }
 
-uint32_t ETHOSU_PMU_Get_IRQ_Enable()
+uint32_t ETHOSU_PMU_Get_IRQ_Enable(void)
 {
-    return read_reg(NPU_REG_PMINTSET);
+    return ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMINTSET);
 }
 
 void ETHOSU_PMU_CNTR_Increment(uint32_t mask)
@@ -233,8 +244,8 @@ void ETHOSU_PMU_CNTR_Increment(uint32_t mask)
         {
             ETHOSU_PMU_CNTR_Disable(ETHOSU_PMU_CCNT_Msk);
             uint64_t val = ETHOSU_PMU_Get_CCNTR() + 1;
-            write_reg(NPU_REG_PMCCNTR_LO, (val & MASK_0_31_BITS));
-            write_reg(NPU_REG_PMCCNTR_HI, (val & MASK_32_47_BITS) >> 32);
+            ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_LO, (val & MASK_0_31_BITS));
+            ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_HI, (val & MASK_32_47_BITS) >> 32);
             if (cntrs_active & ETHOSU_PMU_CCNT_Msk)
             {
                 ETHOSU_PMU_CNTR_Enable(ETHOSU_PMU_CCNT_Msk);
@@ -248,8 +259,8 @@ void ETHOSU_PMU_CNTR_Increment(uint32_t mask)
         if (mask & cntr)
         {
             ETHOSU_PMU_CNTR_Disable(cntr);
-            uint32_t val = read_reg(NPU_REG_PMEVCNTR(i));
-            write_reg(NPU_REG_PMEVCNTR(i), val + 1);
+            uint32_t val = ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMEVCNTR(i));
+            ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMEVCNTR(i), val + 1);
             if (cntrs_active & cntr)
             {
                 ETHOSU_PMU_CNTR_Enable(cntr);
@@ -261,15 +272,15 @@ void ETHOSU_PMU_CNTR_Increment(uint32_t mask)
 void ETHOSU_PMU_PMCCNTR_CFG_Set_Start_Event(uint32_t start_event)
 {
     struct pmccntr_cfg_r cfg;
-    cfg.word                = read_reg(NPU_REG_PMCCNTR_CFG);
+    cfg.word                = ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_CFG);
     cfg.CYCLE_CNT_CFG_START = start_event & ETHOSU_PMCCNTR_CFG_START_STOP_EVENT_MASK;
-    write_reg(NPU_REG_PMCCNTR_CFG, cfg.word);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_CFG, cfg.word);
 }
 
 void ETHOSU_PMU_PMCCNTR_CFG_Set_Stop_Event(uint32_t stop_event)
 {
     struct pmccntr_cfg_r cfg;
-    cfg.word               = read_reg(NPU_REG_PMCCNTR_CFG);
+    cfg.word               = ethosu_read_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_CFG);
     cfg.CYCLE_CNT_CFG_STOP = stop_event & ETHOSU_PMCCNTR_CFG_START_STOP_EVENT_MASK;
-    write_reg(NPU_REG_PMCCNTR_CFG, cfg.word);
+    ethosu_write_reg(&ethosu_drv.dev, NPU_REG_PMCCNTR_CFG, cfg.word);
 }