1 files changed, 16 insertions, 5 deletions

diff --git a/drivers/net/mlx4/icm.c b/drivers/net/mlx4/icm.c
index 250e24887578..4b3c109d5eae 100644
--- a/drivers/net/mlx4/icm.c
+++ b/drivers/net/mlx4/icm.c
@@ -301,9 +301,9 @@ void mlx4_table_put(struct mlx4_dev *dev, struct mlx4_icm_table *table, int obj)
         mutex_unlock(&table->mutex);
 }
 
-void *mlx4_table_find(struct mlx4_icm_table *table, int obj)
+void *mlx4_table_find(struct mlx4_icm_table *table, int obj, dma_addr_t *dma_handle)
 {
-        int idx, offset, i;
+        int idx, offset, dma_offset, i;
         struct mlx4_icm_chunk *chunk;
         struct mlx4_icm *icm;
         struct page *page = NULL;
@@ -313,15 +313,26 @@ void *mlx4_table_find(struct mlx4_icm_table *table, int obj)
 
         mutex_lock(&table->mutex);
 
-        idx = obj & (table->num_obj - 1);
-        icm = table->icm[idx / (MLX4_TABLE_CHUNK_SIZE / table->obj_size)];
-        offset = idx % (MLX4_TABLE_CHUNK_SIZE / table->obj_size);
+        idx = (obj & (table->num_obj - 1)) * table->obj_size;
+        icm = table->icm[idx / MLX4_TABLE_CHUNK_SIZE];
+        dma_offset = offset = idx % MLX4_TABLE_CHUNK_SIZE;
 
         if (!icm)
                 goto out;
 
         list_for_each_entry(chunk, &icm->chunk_list, list) {
                 for (i = 0; i < chunk->npages; ++i) {
+                        if (dma_handle && dma_offset >= 0) {
+                                if (sg_dma_len(&chunk->mem[i]) > dma_offset)
+                                        *dma_handle = sg_dma_address(&chunk->mem[i]) +
+                                                dma_offset;
+                                dma_offset -= sg_dma_len(&chunk->mem[i]);
+                        }
+                        /*
+                         * DMA mapping can merge pages but not split them,
+                         * so if we found the page, dma_handle has already
+                         * been assigned to.
+                         */
                         if (chunk->mem[i].length > offset) {
                                 page = chunk->mem[i].page;
                                 goto out;