diff --git a/drivers/media/platform/chips-media/wave5/wave5-hw.c b/drivers/media/platform/chips-media/wave5/wave5-hw.c
index cdd0a0948a94..36f2fc818013 100644
--- a/drivers/media/platform/chips-media/wave5/wave5-hw.c
+++ b/drivers/media/platform/chips-media/wave5/wave5-hw.c
@@ -843,6 +843,36 @@ int wave5_vpu_dec_register_framebuffer(struct vpu_instance *inst, struct frame_b
 	return ret;
 }
 
+static u32 wave5_vpu_dec_validate_sec_axi(struct vpu_instance *inst)
+{
+	struct dec_info *p_dec_info = &inst->codec_info->dec_info;
+	u32 bit_size = 0, ip_size = 0, lf_size = 0, ret = 0;
+	u32 sram_size = inst->dev->sram_size;
+
+	if (!sram_size)
+		return 0;
+
+	/*
+	 * TODO: calculate bit_size, ip_size, lf_size from inst->src_fmt.width
+	 * and inst->codec_info->dec_info.initial_info.luma_bitdepth
+	 */
+
+	if (p_dec_info->sec_axi_info.use_bit_enable && sram_size >= bit_size) {
+		ret |= BIT(0);
+		sram_size -= bit_size;
+	}
+
+	if (p_dec_info->sec_axi_info.use_ip_enable && sram_size >= ip_size) {
+		ret |= BIT(9);
+		sram_size -= ip_size;
+	}
+
+	if (p_dec_info->sec_axi_info.use_lf_row_enable && sram_size >= lf_size)
+		ret |= BIT(15);
+
+	return ret;
+}
+
 int wave5_vpu_decode(struct vpu_instance *inst, u32 *fail_res)
 {
 	u32 reg_val;
@@ -855,9 +885,7 @@ int wave5_vpu_decode(struct vpu_instance *inst, u32 *fail_res)
 	vpu_write_reg(inst->dev, W5_BS_OPTION, get_bitstream_options(p_dec_info));
 
 	/* secondary AXI */
-	reg_val = p_dec_info->sec_axi_info.use_bit_enable |
-		(p_dec_info->sec_axi_info.use_ip_enable << 9) |
-		(p_dec_info->sec_axi_info.use_lf_row_enable << 15);
+	reg_val = wave5_vpu_dec_validate_sec_axi(inst);
 	vpu_write_reg(inst->dev, W5_USE_SEC_AXI, reg_val);
 
 	/* set attributes of user buffer */
@@ -1938,6 +1966,31 @@ int wave5_vpu_enc_register_framebuffer(struct device *dev, struct vpu_instance *
 	return ret;
 }
 
+static u32 wave5_vpu_enc_validate_sec_axi(struct vpu_instance *inst)
+{
+	struct enc_info *p_enc_info = &inst->codec_info->enc_info;
+	u32 rdo_size = 0, lf_size = 0, ret = 0;
+	u32 sram_size = inst->dev->sram_size;
+
+	if (!sram_size)
+		return 0;
+
+	/*
+	 * TODO: calculate rdo_size and lf_size from inst->src_fmt.width and
+	 * inst->codec_info->enc_info.open_param.wave_param.internal_bit_depth
+	 */
+
+	if (p_enc_info->sec_axi_info.use_enc_rdo_enable && sram_size >= rdo_size) {
+		ret |= BIT(11);
+		sram_size -= rdo_size;
+	}
+
+	if (p_enc_info->sec_axi_info.use_enc_lf_enable && sram_size >= lf_size)
+		ret |= BIT(15);
+
+	return ret;
+}
+
 int wave5_vpu_encode(struct vpu_instance *inst, struct enc_param *option, u32 *fail_res)
 {
 	u32 src_frame_format;
@@ -1959,8 +2012,7 @@ int wave5_vpu_encode(struct vpu_instance *inst, struct enc_param *option, u32 *f
 
 	vpu_write_reg(inst->dev, W5_CMD_ENC_PIC_SRC_AXI_SEL, DEFAULT_SRC_AXI);
 	/* secondary AXI */
-	reg_val = (p_enc_info->sec_axi_info.use_enc_rdo_enable << 11) |
-		(p_enc_info->sec_axi_info.use_enc_lf_enable << 15);
+	reg_val = wave5_vpu_enc_validate_sec_axi(inst);
 	vpu_write_reg(inst->dev, W5_CMD_ENC_PIC_USE_SEC_AXI, reg_val);
 
 	vpu_write_reg(inst->dev, W5_CMD_ENC_PIC_REPORT_PARAM, 0);
diff --git a/drivers/media/platform/chips-media/wave5/wave5-vdi.c b/drivers/media/platform/chips-media/wave5/wave5-vdi.c
index 3809f70bc0b4..556de2f043fe 100644
--- a/drivers/media/platform/chips-media/wave5/wave5-vdi.c
+++ b/drivers/media/platform/chips-media/wave5/wave5-vdi.c
@@ -174,16 +174,19 @@ int wave5_vdi_allocate_array(struct vpu_device *vpu_dev, struct vpu_buf *array,
 void wave5_vdi_allocate_sram(struct vpu_device *vpu_dev)
 {
 	struct vpu_buf *vb = &vpu_dev->sram_buf;
+	dma_addr_t daddr;
+	void *vaddr;
+	size_t size;
 
-	if (!vpu_dev->sram_pool || !vpu_dev->sram_size)
+	if (!vpu_dev->sram_pool || vb->vaddr)
 		return;
 
-	if (!vb->vaddr) {
-		vb->size = vpu_dev->sram_size;
-		vb->vaddr = gen_pool_dma_alloc(vpu_dev->sram_pool, vb->size,
-					       &vb->daddr);
-		if (!vb->vaddr)
-			vb->size = 0;
+	size = min_t(size_t, vpu_dev->sram_size, gen_pool_avail(vpu_dev->sram_pool));
+	vaddr = gen_pool_dma_alloc(vpu_dev->sram_pool, size, &daddr);
+	if (vaddr) {
+		vb->vaddr = vaddr;
+		vb->daddr = daddr;
+		vb->size = size;
 	}
 
 	dev_dbg(vpu_dev->dev, "%s: sram daddr: %pad, size: %zu, vaddr: 0x%p\n",
@@ -197,9 +200,7 @@ void wave5_vdi_free_sram(struct vpu_device *vpu_dev)
 	if (!vb->size || !vb->vaddr)
 		return;
 
-	if (vb->vaddr)
-		gen_pool_free(vpu_dev->sram_pool, (unsigned long)vb->vaddr,
-			      vb->size);
+	gen_pool_free(vpu_dev->sram_pool, (unsigned long)vb->vaddr, vb->size);
 
 	memset(vb, 0, sizeof(*vb));
 }
diff --git a/drivers/media/platform/chips-media/wave5/wave5-vpu.c b/drivers/media/platform/chips-media/wave5/wave5-vpu.c
index 1e631da58e15..9e93969ab6db 100644
--- a/drivers/media/platform/chips-media/wave5/wave5-vpu.c
+++ b/drivers/media/platform/chips-media/wave5/wave5-vpu.c
@@ -25,6 +25,7 @@
 struct wave5_match_data {
 	int flags;
 	const char *fw_name;
+	u32 sram_size;
 };
 
 int wave5_vpu_wait_interrupt(struct vpu_instance *inst, unsigned int timeout)
@@ -177,17 +178,12 @@ static int wave5_vpu_probe(struct platform_device *pdev)
 		goto err_reset_assert;
 	}
 
-	ret = of_property_read_u32(pdev->dev.of_node, "sram-size",
-				   &dev->sram_size);
-	if (ret) {
-		dev_warn(&pdev->dev, "sram-size not found\n");
-		dev->sram_size = 0;
-	}
-
 	dev->sram_pool = of_gen_pool_get(pdev->dev.of_node, "sram", 0);
 	if (!dev->sram_pool)
 		dev_warn(&pdev->dev, "sram node not found\n");
 
+	dev->sram_size = match_data->sram_size;
+
 	dev->product_code = wave5_vdi_read_register(dev, VPU_PRODUCT_CODE_REGISTER);
 	ret = wave5_vdi_init(&pdev->dev);
 	if (ret < 0) {
@@ -281,6 +277,7 @@ static void wave5_vpu_remove(struct platform_device *pdev)
 static const struct wave5_match_data ti_wave521c_data = {
 	.flags = WAVE5_IS_ENC | WAVE5_IS_DEC,
 	.fw_name = "cnm/wave521c_k3_codec_fw.bin",
+	.sram_size = (64 * 1024),
 };
 
 static const struct of_device_id wave5_dt_ids[] = {
