diff --git a/drivers/gpu/nvgpu/common/gr/gr_falcon.c b/drivers/gpu/nvgpu/common/gr/gr_falcon.c
index 729d648b7..d48049919 100644
--- a/drivers/gpu/nvgpu/common/gr/gr_falcon.c
+++ b/drivers/gpu/nvgpu/common/gr/gr_falcon.c
@@ -237,9 +237,12 @@ static int nvgpu_gr_falcon_init_ctxsw_ucode_vaspace(struct gk20a *g,
 static void nvgpu_gr_falcon_init_ctxsw_ucode_segment(
 	struct nvgpu_ctxsw_ucode_segment *p_seg, u32 *offset, u32 size)
 {
+	u32 ucode_offset;
+
 	p_seg->offset = *offset;
 	p_seg->size = size;
-	*offset = ALIGN(*offset + size, SZ_256);
+	ucode_offset = nvgpu_safe_add_u32(*offset, size);
+	*offset = ALIGN(ucode_offset, 256U);
 }
 
 static void nvgpu_gr_falcon_init_ctxsw_ucode_segments(