diff --git a/src/egl/drivers/dri2/egl_dri2.c b/src/egl/drivers/dri2/egl_dri2.c
index f68a15c8a21..66d3789a0b9 100644
--- a/src/egl/drivers/dri2/egl_dri2.c
+++ b/src/egl/drivers/dri2/egl_dri2.c
@@ -796,7 +796,7 @@ dri2_setup_screen(_EGLDisplay *disp)
    if (dri2_dpy->flush_control)
       disp->Extensions.KHR_context_flush_control = EGL_TRUE;
 
-   if (dri2_dpy->buffer_damage)
+   if (get_pipe_screen(disp)->set_damage_region)
       disp->Extensions.KHR_partial_update = EGL_TRUE;
 
    disp->Extensions.EXT_protected_surface =
@@ -1756,7 +1756,7 @@ dri2_swap_buffers(_EGLDisplay *disp, _EGLSurface *surf)
    /* SwapBuffers marks the end of the frame; reset the damage region for
     * use again next time.
     */
-   if (ret && dri2_dpy->buffer_damage)
+   if (ret && disp->Extensions.KHR_partial_update)
       dri2_dpy->buffer_damage->set_damage_region(dri_drawable, 0, NULL);
 
    return ret;
@@ -1782,7 +1782,7 @@ dri2_swap_buffers_with_damage(_EGLDisplay *disp, _EGLSurface *surf,
    /* SwapBuffers marks the end of the frame; reset the damage region for
     * use again next time.
     */
-   if (ret && dri2_dpy->buffer_damage)
+   if (ret && disp->Extensions.KHR_partial_update)
       dri2_dpy->buffer_damage->set_damage_region(dri_drawable, 0, NULL);
 
    return ret;
@@ -1803,7 +1803,7 @@ dri2_swap_buffers_region(_EGLDisplay *disp, _EGLSurface *surf, EGLint numRects,
    /* SwapBuffers marks the end of the frame; reset the damage region for
     * use again next time.
     */
-   if (ret && dri2_dpy->buffer_damage)
+   if (ret && disp->Extensions.KHR_partial_update)
       dri2_dpy->buffer_damage->set_damage_region(dri_drawable, 0, NULL);
 
    return ret;
@@ -1816,7 +1816,7 @@ dri2_set_damage_region(_EGLDisplay *disp, _EGLSurface *surf, EGLint *rects,
    struct dri2_egl_display *dri2_dpy = dri2_egl_display_lock(disp);
    __DRIdrawable *drawable = dri2_dpy->vtbl->get_dri_drawable(surf);
 
-   if (!dri2_dpy->buffer_damage) {
+   if (!disp->Extensions.KHR_partial_update) {
       mtx_unlock(&dri2_dpy->lock);
       return EGL_FALSE;
    }