diff --git a/libdmabufheap/include/dmabuf_alloc.h b/libdmabufheap/include/dmabuf_alloc.h
index e8233198269f8a818663f971f00135d625e90ead..abec18f59d6c56f8330e91033ba40670a17aa049 100644
--- a/libdmabufheap/include/dmabuf_alloc.h
+++ b/libdmabufheap/include/dmabuf_alloc.h
@@ -38,9 +38,18 @@ typedef enum {
 typedef struct {
     unsigned int fd;
     size_t size;
-    unsigned int heapFlags;
+    __u64 heapFlags;
 } DmabufHeapBuffer;
 
+enum DmaHeapFlagOwnerId {
+    DMA_OWNER_DEFAULT,
+    DMA_OWNER_GPU,
+    DMA_OWNER_MEDIA_CODEC,
+    COUNT_DMA_OWNER,
+};
+
+void SetOwnerIdForHeapFlags(DmabufHeapBuffer *buffer, enum DmaHeapFlagOwnerId ownerId);
+
 int DmabufHeapOpen(const char *heapName);
 
 int DmabufHeapClose(unsigned int fd);
diff --git a/libdmabufheap/src/dmabuf_alloc.c b/libdmabufheap/src/dmabuf_alloc.c
index b283b78b6f7a47b80f4cb110a89a012d7e1a1c49..79e7d539e560f75905204ebe2ceef9be026baff8 100644
--- a/libdmabufheap/src/dmabuf_alloc.c
+++ b/libdmabufheap/src/dmabuf_alloc.c
@@ -51,6 +51,13 @@ static bool IsSyncTypeValid(DmabufHeapBufferSyncType syncType)
     }
 }
 
+void SetOwnerIdForHeapFlags(DmabufHeapBuffer *buffer, enum DmaHeapFlagOwnerId ownerId)
+{
+    if (buffer) {
+        set_owner_id_for_heap_flags(&buffer->heapFlags, ownerId);
+    }
+}
+
 int DmabufHeapOpen(const char *heapName)
 {
     if (!IsHeapNameValid(heapName)) {
diff --git a/libdmabufheap/test/unittest/libdmabufheap/dmabuf_alloc_test.cpp b/libdmabufheap/test/unittest/libdmabufheap/dmabuf_alloc_test.cpp
index 9a318fce930a9e5e8bad325f21f84ecc18ef7276..1ad1e5680387294c4cfc1ffc6c00822349208a9a 100644
--- a/libdmabufheap/test/unittest/libdmabufheap/dmabuf_alloc_test.cpp
+++ b/libdmabufheap/test/unittest/libdmabufheap/dmabuf_alloc_test.cpp
@@ -83,6 +83,13 @@ HWTEST_F(DmabufAllocTest, AllocSingleBuffer, TestSize.Level1)
     ASSERT_GE(heapFd, 0);
 
     DmabufHeapBuffer buffer = { .size = BUFFER_SIZE, .heapFlags = 0 };
+
+    SetOwnerIdForHeapFlags(&buffer, DMA_OWNER_MEDIA_CODEC);
+
+    __u64 ownerId = get_owner_id_from_heap_flags(buffer.heapFlags);
+
+    ASSERT_EQ(DMA_OWNER_MEDIA_CODEC, ownerId);
+
     ASSERT_EQ(0, DmabufHeapBufferAlloc(heapFd, &buffer));
 
     void *ptr = mmap(NULL, BUFFER_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, buffer.fd, 0);