diff --git a/lib/sw_apps/openamp_matrix_multiply/src/baremetal.c b/lib/sw_apps/openamp_matrix_multiply/src/baremetal.c
index 2e559fe1..2b225aae 100644
--- a/lib/sw_apps/openamp_matrix_multiply/src/baremetal.c
+++ b/lib/sw_apps/openamp_matrix_multiply/src/baremetal.c
@@ -36,6 +36,7 @@
 #include "xil_cache.h"
 #include "xil_mmu.h"
 #include "baremetal.h"
+#include "env.h"
 
 XScuGic InterruptController;
 
@@ -87,38 +88,6 @@ void zynqMP_r5_irq_isr() {
 	XScuGic_CPUWriteReg(&InterruptController,XSCUGIC_EOI_OFFSET, raw_irq);
 }
 
-/***********************************************************************
- *
- *
- * zynqMP_r5_map_mem_region
- *
- *
- * This function sets-up the region of memory based on the given
- * attributes
- * There is no MMU for R5, no need to map phy address to vrt_addr
- *
- * @param addr		     - Starting address of memory region
- * @parma size           - size of region
- * @param attrib  		 - Attributes for memory region
- *
- *
- * OUTPUTS
- *
- *       None
- *
- ***********************************************************************/
-void zynqMP_r5_map_mem_region(u32 addr, u32 size, u32 attrib) {
-	u32 Index, NumSize;
-	/* Calculating the number of MBs required for the shared region*/
-	NumSize = size / 0x100000;
-
-	/* Xil_SetTlbAttributes is designed to configure memory for 1MB
-	 * region. The API is called multiple times to configure the number
-	 * of MBs required by shared memory size (calculated as NumSize)*/
-	for (Index = 0; Index < NumSize; Index ++)
-		Xil_SetTlbAttributes(addr + 0x100000 * Index, attrib);
-}
-
 /*
  ***********************************************************************
  * IPI handling
@@ -230,6 +199,28 @@ void platform_cache_disable() {
 }
 
 void platform_map_mem_region(unsigned int va,unsigned int pa, unsigned int size,unsigned int flags) {
+
+	unsigned int r5_flags;
+
+	/* Assume DEVICE_SHARED if nothing indicates this is memory.  */
+	r5_flags = DEVICE_SHARED;
+	if (flags & SHARED_MEM) {
+		r5_flags = NORM_SHARED_NCACHE;
+		if (flags & WB_CACHE) {
+			r5_flags = NORM_SHARED_WB_WA;
+		} else if (flags & WT_CACHE) {
+			r5_flags = NORM_SHARED_WT_NWA;
+		}
+	} else if (flags & MEM_MAPPED) {
+		r5_flags = NORM_NSHARED_NCACHE;
+		if (flags & WB_CACHE) {
+			r5_flags = NORM_NSHARED_WB_WA;
+		} else if (flags & WT_CACHE) {
+			r5_flags = NORM_NSHARED_WT_NWA;
+		}
+	}
+
+	Xil_SetMPURegion(pa, size, r5_flags | PRIV_RW_USER_RW);
 	return;
 }
 
diff --git a/lib/sw_apps/openamp_matrix_multiply/src/matrix_multiply.c b/lib/sw_apps/openamp_matrix_multiply/src/matrix_multiply.c
index 029c0953..c58dd5de 100644
--- a/lib/sw_apps/openamp_matrix_multiply/src/matrix_multiply.c
+++ b/lib/sw_apps/openamp_matrix_multiply/src/matrix_multiply.c
@@ -88,14 +88,6 @@
 #define NUM_MATRIX      2
 #define SHUTDOWN_MSG    0xEF56A55A
 
-/*
- * Shared memory location as defined in linux device tree for remoteproc
- * User may need to check with device tree of remoteproc and ensure the
- * share memory address is same
- */
-#define SHARED_MEMORY 	0x3ED00000
-#define SHARED_SIZE 	0x400000	/* size of the shared memory*/
-
 typedef struct _matrix {
 	unsigned int size;
 	unsigned int elements[MAX_SIZE][MAX_SIZE];
@@ -184,9 +176,6 @@ static void Matrix_Multiply(const matrix *m, const matrix *n, matrix *r) {
 }
 
 static void init_system() {
-	/* configure MPU for shared memory region */
-	zynqMP_r5_map_mem_region(SHARED_MEMORY, SHARED_SIZE, NORM_SHARED_NCACHE | PRIV_RW_USER_RW);
-
 
 	/* Initilaize GIC */
 	zynqMP_r5_gic_initialize();