int
 i915_gem_object_get_pages_buddy(struct drm_i915_gem_object *obj)
 {
+       const u64 max_segment = i915_sg_segment_size();
        struct intel_memory_region *mem = obj->mm.region;
        struct list_head *blocks = &obj->mm.blocks;
        resource_size_t size = obj->base.size;
        if (!st)
                return -ENOMEM;
 
-       if (sg_alloc_table(st, size >> ilog2(mem->mm.chunk_size), GFP_KERNEL)) {
+       if (sg_alloc_table(st, size >> PAGE_SHIFT, GFP_KERNEL)) {
                kfree(st);
                return -ENOMEM;
        }
                                   i915_buddy_block_size(&mem->mm, block));
                offset = i915_buddy_block_offset(block);
 
-               GEM_BUG_ON(overflows_type(block_size, sg->length));
+               while (block_size) {
+                       u64 len;
 
-               if (offset != prev_end ||
-                   add_overflows_t(typeof(sg->length), sg->length, block_size)) {
-                       if (st->nents) {
-                               sg_page_sizes |= sg->length;
-                               sg = __sg_next(sg);
+                       if (offset != prev_end || sg->length >= max_segment) {
+                               if (st->nents) {
+                                       sg_page_sizes |= sg->length;
+                                       sg = __sg_next(sg);
+                               }
+
+                               sg_dma_address(sg) = mem->region.start + offset;
+                               sg_dma_len(sg) = 0;
+                               sg->length = 0;
+                               st->nents++;
                        }
 
-                       sg_dma_address(sg) = mem->region.start + offset;
-                       sg_dma_len(sg) = block_size;
+                       len = min(block_size, max_segment - sg->length);
+                       sg->length += len;
+                       sg_dma_len(sg) += len;
 
-                       sg->length = block_size;
+                       offset += len;
+                       block_size -= len;
 
-                       st->nents++;
-               } else {
-                       sg->length += block_size;
-                       sg_dma_len(sg) += block_size;
+                       prev_end = offset;
                }
-
-               prev_end = offset + block_size;
        }
 
        sg_page_sizes |= sg->length;
 
        i915_gem_object_put(obj);
 }
 
+static bool is_contiguous(struct drm_i915_gem_object *obj)
+{
+       struct scatterlist *sg;
+       dma_addr_t addr = -1;
+
+       for (sg = obj->mm.pages->sgl; sg; sg = sg_next(sg)) {
+               if (addr != -1 && sg_dma_address(sg) != addr)
+                       return false;
+
+               addr = sg_dma_address(sg) + sg_dma_len(sg);
+       }
+
+       return true;
+}
+
 static int igt_mock_contiguous(void *arg)
 {
        struct intel_memory_region *mem = arg;
        if (IS_ERR(obj))
                return PTR_ERR(obj);
 
-       if (obj->mm.pages->nents != 1) {
-               pr_err("%s min object spans multiple sg entries\n", __func__);
+       if (!is_contiguous(obj)) {
+               pr_err("%s min object spans disjoint sg entries\n", __func__);
                err = -EINVAL;
                goto err_close_objects;
        }
        if (IS_ERR(obj))
                return PTR_ERR(obj);
 
-       if (obj->mm.pages->nents != 1) {
-               pr_err("%s max object spans multiple sg entries\n", __func__);
+       if (!is_contiguous(obj)) {
+               pr_err("%s max object spans disjoint sg entries\n", __func__);
                err = -EINVAL;
                goto err_close_objects;
        }
                goto err_close_objects;
        }
 
-       if (obj->mm.pages->nents != 1) {
-               pr_err("%s object spans multiple sg entries\n", __func__);
+       if (!is_contiguous(obj)) {
+               pr_err("%s object spans disjoint sg entries\n", __func__);
                err = -EINVAL;
                goto err_close_objects;
        }