drivers/misc/habanalabs/common/memory.c

   1 // SPDX-License-Identifier: GPL-2.0
   2
   3 /*
   4  * Copyright 2016-2019 HabanaLabs, Ltd.
   5  * All Rights Reserved.
   6  */
   7
   8 #include <uapi/misc/habanalabs.h>
   9 #include "habanalabs.h"
  10 #include "../include/hw_ip/mmu/mmu_general.h"
  11
  12 #include <linux/uaccess.h>
  13 #include <linux/slab.h>
  14
  15 #define HL_MMU_DEBUG    0
  16
  17 /* use small pages for supporting non-pow2 (32M/40M/48M) DRAM phys page sizes */
  18 #define DRAM_POOL_PAGE_SIZE SZ_8M
  19
  20 /*
  21  * The va ranges in context object contain a list with the available chunks of
  22  * device virtual memory.
  23  * There is one range for host allocations and one for DRAM allocations.
  24  *
  25  * On initialization each range contains one chunk of all of its available
  26  * virtual range which is a half of the total device virtual range.
  27  *
  28  * On each mapping of physical pages, a suitable virtual range chunk (with a
  29  * minimum size) is selected from the list. If the chunk size equals the
  30  * requested size, the chunk is returned. Otherwise, the chunk is split into
  31  * two chunks - one to return as result and a remainder to stay in the list.
  32  *
  33  * On each Unmapping of a virtual address, the relevant virtual chunk is
  34  * returned to the list. The chunk is added to the list and if its edges match
  35  * the edges of the adjacent chunks (means a contiguous chunk can be created),
  36  * the chunks are merged.
  37  *
  38  * On finish, the list is checked to have only one chunk of all the relevant
  39  * virtual range (which is a half of the device total virtual range).
  40  * If not (means not all mappings were unmapped), a warning is printed.
  41  */
  42
  43 /*
  44  * alloc_device_memory() - allocate device memory.
  45  * @ctx: pointer to the context structure.
  46  * @args: host parameters containing the requested size.
  47  * @ret_handle: result handle.
  48  *
  49  * This function does the following:
  50  * - Allocate the requested size rounded up to 'dram_page_size' pages.
  51  * - Return unique handle for later map/unmap/free.
  52  */
  53 static int alloc_device_memory(struct hl_ctx *ctx, struct hl_mem_in *args,
  54                                 u32 *ret_handle)
  55 {
  56         struct hl_device *hdev = ctx->hdev;
  57         struct hl_vm *vm = &hdev->vm;
  58         struct hl_vm_phys_pg_pack *phys_pg_pack;
  59         u64 paddr = 0, total_size, num_pgs, i;
  60         u32 num_curr_pgs, page_size;
  61         int handle, rc;
  62         bool contiguous;
  63
  64         num_curr_pgs = 0;
  65         page_size = hdev->asic_prop.dram_page_size;
  66         num_pgs = DIV_ROUND_UP_ULL(args->alloc.mem_size, page_size);
  67         total_size = num_pgs * page_size;
  68
  69         if (!total_size) {
  70                 dev_err(hdev->dev, "Cannot allocate 0 bytes\n");
  71                 return -EINVAL;
  72         }
  73
  74         contiguous = args->flags & HL_MEM_CONTIGUOUS;
  75
  76         if (contiguous) {
  77                 paddr = (u64) gen_pool_alloc(vm->dram_pg_pool, total_size);
  78                 if (!paddr) {
  79                         dev_err(hdev->dev,
  80                                 "failed to allocate %llu contiguous pages with total size of %llu\n",
  81                                 num_pgs, total_size);
  82                         return -ENOMEM;
  83                 }
  84
  85                 if (hdev->memory_scrub) {
  86                         rc = hdev->asic_funcs->scrub_device_mem(hdev, paddr,
  87                                         total_size);
  88                         if (rc) {
  89                                 dev_err(hdev->dev,
  90                                         "Failed to scrub contiguous device memory\n");
  91                                 goto pages_pack_err;
  92                         }
  93                 }
  94         }
  95
  96         phys_pg_pack = kzalloc(sizeof(*phys_pg_pack), GFP_KERNEL);
  97         if (!phys_pg_pack) {
  98                 rc = -ENOMEM;
  99                 goto pages_pack_err;
 100         }
 101
 102         phys_pg_pack->vm_type = VM_TYPE_PHYS_PACK;
 103         phys_pg_pack->asid = ctx->asid;
 104         phys_pg_pack->npages = num_pgs;
 105         phys_pg_pack->page_size = page_size;
 106         phys_pg_pack->total_size = total_size;
 107         phys_pg_pack->flags = args->flags;
 108         phys_pg_pack->contiguous = contiguous;
 109
 110         phys_pg_pack->pages = kvmalloc_array(num_pgs, sizeof(u64), GFP_KERNEL);
 111         if (ZERO_OR_NULL_PTR(phys_pg_pack->pages)) {
 112                 rc = -ENOMEM;
 113                 goto pages_arr_err;
 114         }
 115
 116         if (phys_pg_pack->contiguous) {
 117                 for (i = 0 ; i < num_pgs ; i++)
 118                         phys_pg_pack->pages[i] = paddr + i * page_size;
 119         } else {
 120                 for (i = 0 ; i < num_pgs ; i++) {
 121                         phys_pg_pack->pages[i] = (u64) gen_pool_alloc(
 122                                                         vm->dram_pg_pool,
 123                                                         page_size);
 124                         if (!phys_pg_pack->pages[i]) {
 125                                 dev_err(hdev->dev,
 126                                         "Failed to allocate device memory (out of memory)\n");
 127                                 rc = -ENOMEM;
 128                                 goto page_err;
 129                         }
 130
 131                         if (hdev->memory_scrub) {
 132                                 rc = hdev->asic_funcs->scrub_device_mem(hdev,
 133                                                 phys_pg_pack->pages[i],
 134                                                 page_size);
 135                                 if (rc) {
 136                                         dev_err(hdev->dev,
 137                                                 "Failed to scrub device memory\n");
 138                                         goto page_err;
 139                                 }
 140                         }
 141
 142                         num_curr_pgs++;
 143                 }
 144         }
 145
 146         spin_lock(&vm->idr_lock);
 147         handle = idr_alloc(&vm->phys_pg_pack_handles, phys_pg_pack, 1, 0,
 148                                 GFP_ATOMIC);
 149         spin_unlock(&vm->idr_lock);
 150
 151         if (handle < 0) {
 152                 dev_err(hdev->dev, "Failed to get handle for page\n");
 153                 rc = -EFAULT;
 154                 goto idr_err;
 155         }
 156
 157         for (i = 0 ; i < num_pgs ; i++)
 158                 kref_get(&vm->dram_pg_pool_refcount);
 159
 160         phys_pg_pack->handle = handle;
 161
 162         atomic64_add(phys_pg_pack->total_size, &ctx->dram_phys_mem);
 163         atomic64_add(phys_pg_pack->total_size, &hdev->dram_used_mem);
 164
 165         *ret_handle = handle;
 166
 167         return 0;
 168
 169 idr_err:
 170 page_err:
 171         if (!phys_pg_pack->contiguous)
 172                 for (i = 0 ; i < num_curr_pgs ; i++)
 173                         gen_pool_free(vm->dram_pg_pool, phys_pg_pack->pages[i],
 174                                         page_size);
 175
 176         kvfree(phys_pg_pack->pages);
 177 pages_arr_err:
 178         kfree(phys_pg_pack);
 179 pages_pack_err:
 180         if (contiguous)
 181                 gen_pool_free(vm->dram_pg_pool, paddr, total_size);
 182
 183         return rc;
 184 }
 185
 186 /**
 187  * dma_map_host_va() - DMA mapping of the given host virtual address.
 188  * @hdev: habanalabs device structure.
 189  * @addr: the host virtual address of the memory area.
 190  * @size: the size of the memory area.
 191  * @p_userptr: pointer to result userptr structure.
 192  *
 193  * This function does the following:
 194  * - Allocate userptr structure.
 195  * - Pin the given host memory using the userptr structure.
 196  * - Perform DMA mapping to have the DMA addresses of the pages.
 197  */
 198 static int dma_map_host_va(struct hl_device *hdev, u64 addr, u64 size,
 199                                 struct hl_userptr **p_userptr)
 200 {
 201         struct hl_userptr *userptr;
 202         int rc;
 203
 204         userptr = kzalloc(sizeof(*userptr), GFP_KERNEL);
 205         if (!userptr) {
 206                 rc = -ENOMEM;
 207                 goto userptr_err;
 208         }
 209
 210         rc = hl_pin_host_memory(hdev, addr, size, userptr);
 211         if (rc) {
 212                 dev_err(hdev->dev, "Failed to pin host memory\n");
 213                 goto pin_err;
 214         }
 215
 216         rc = hdev->asic_funcs->asic_dma_map_sg(hdev, userptr->sgt->sgl,
 217                                         userptr->sgt->nents, DMA_BIDIRECTIONAL);
 218         if (rc) {
 219                 dev_err(hdev->dev, "failed to map sgt with DMA region\n");
 220                 goto dma_map_err;
 221         }
 222
 223         userptr->dma_mapped = true;
 224         userptr->dir = DMA_BIDIRECTIONAL;
 225         userptr->vm_type = VM_TYPE_USERPTR;
 226
 227         *p_userptr = userptr;
 228
 229         return 0;
 230
 231 dma_map_err:
 232         hl_unpin_host_memory(hdev, userptr);
 233 pin_err:
 234         kfree(userptr);
 235 userptr_err:
 236
 237         return rc;
 238 }
 239
 240 /**
 241  * dma_unmap_host_va() - DMA unmapping of the given host virtual address.
 242  * @hdev: habanalabs device structure.
 243  * @userptr: userptr to free.
 244  *
 245  * This function does the following:
 246  * - Unpins the physical pages.
 247  * - Frees the userptr structure.
 248  */
 249 static void dma_unmap_host_va(struct hl_device *hdev,
 250                                 struct hl_userptr *userptr)
 251 {
 252         hl_unpin_host_memory(hdev, userptr);
 253         kfree(userptr);
 254 }
 255
 256 /**
 257  * dram_pg_pool_do_release() - free DRAM pages pool
 258  * @ref: pointer to reference object.
 259  *
 260  * This function does the following:
 261  * - Frees the idr structure of physical pages handles.
 262  * - Frees the generic pool of DRAM physical pages.
 263  */
 264 static void dram_pg_pool_do_release(struct kref *ref)
 265 {
 266         struct hl_vm *vm = container_of(ref, struct hl_vm,
 267                         dram_pg_pool_refcount);
 268
 269         /*
 270          * free the idr here as only here we know for sure that there are no
 271          * allocated physical pages and hence there are no handles in use
 272          */
 273         idr_destroy(&vm->phys_pg_pack_handles);
 274         gen_pool_destroy(vm->dram_pg_pool);
 275 }
 276
 277 /**
 278  * free_phys_pg_pack() - free physical page pack.
 279  * @hdev: habanalabs device structure.
 280  * @phys_pg_pack: physical page pack to free.
 281  *
 282  * This function does the following:
 283  * - For DRAM memory only, iterate over the pack and free each physical block
 284  *   structure by returning it to the general pool.
 285  * - Free the hl_vm_phys_pg_pack structure.
 286  */
 287 static void free_phys_pg_pack(struct hl_device *hdev,
 288                                 struct hl_vm_phys_pg_pack *phys_pg_pack)
 289 {
 290         struct hl_vm *vm = &hdev->vm;
 291         u64 i;
 292
 293         if (!phys_pg_pack->created_from_userptr) {
 294                 if (phys_pg_pack->contiguous) {
 295                         gen_pool_free(vm->dram_pg_pool, phys_pg_pack->pages[0],
 296                                         phys_pg_pack->total_size);
 297
 298                         for (i = 0; i < phys_pg_pack->npages ; i++)
 299                                 kref_put(&vm->dram_pg_pool_refcount,
 300                                         dram_pg_pool_do_release);
 301                 } else {
 302                         for (i = 0 ; i < phys_pg_pack->npages ; i++) {
 303                                 gen_pool_free(vm->dram_pg_pool,
 304                                                 phys_pg_pack->pages[i],
 305                                                 phys_pg_pack->page_size);
 306                                 kref_put(&vm->dram_pg_pool_refcount,
 307                                         dram_pg_pool_do_release);
 308                         }
 309                 }
 310         }
 311
 312         kvfree(phys_pg_pack->pages);
 313         kfree(phys_pg_pack);
 314 }
 315
 316 /**
 317  * free_device_memory() - free device memory.
 318  * @ctx: pointer to the context structure.
 319  * @args: host parameters containing the requested size.
 320  *
 321  * This function does the following:
 322  * - Free the device memory related to the given handle.
 323  */
 324 static int free_device_memory(struct hl_ctx *ctx, struct hl_mem_in *args)
 325 {
 326         struct hl_device *hdev = ctx->hdev;
 327         struct hl_vm *vm = &hdev->vm;
 328         struct hl_vm_phys_pg_pack *phys_pg_pack;
 329         u32 handle = args->free.handle;
 330
 331         spin_lock(&vm->idr_lock);
 332         phys_pg_pack = idr_find(&vm->phys_pg_pack_handles, handle);
 333         if (phys_pg_pack) {
 334                 if (atomic_read(&phys_pg_pack->mapping_cnt) > 0) {
 335                         dev_err(hdev->dev, "handle %u is mapped, cannot free\n",
 336                                 handle);
 337                         spin_unlock(&vm->idr_lock);
 338                         return -EINVAL;
 339                 }
 340
 341                 /*
 342                  * must remove from idr before the freeing of the physical
 343                  * pages as the refcount of the pool is also the trigger of the
 344                  * idr destroy
 345                  */
 346                 idr_remove(&vm->phys_pg_pack_handles, handle);
 347                 spin_unlock(&vm->idr_lock);
 348
 349                 atomic64_sub(phys_pg_pack->total_size, &ctx->dram_phys_mem);
 350                 atomic64_sub(phys_pg_pack->total_size, &hdev->dram_used_mem);
 351
 352                 free_phys_pg_pack(hdev, phys_pg_pack);
 353         } else {
 354                 spin_unlock(&vm->idr_lock);
 355                 dev_err(hdev->dev,
 356                         "free device memory failed, no match for handle %u\n",
 357                         handle);
 358                 return -EINVAL;
 359         }
 360
 361         return 0;
 362 }
 363
 364 /**
 365  * clear_va_list_locked() - free virtual addresses list.
 366  * @hdev: habanalabs device structure.
 367  * @va_list: list of virtual addresses to free.
 368  *
 369  * This function does the following:
 370  * - Iterate over the list and free each virtual addresses block.
 371  *
 372  * This function should be called only when va_list lock is taken.
 373  */
 374 static void clear_va_list_locked(struct hl_device *hdev,
 375                 struct list_head *va_list)
 376 {
 377         struct hl_vm_va_block *va_block, *tmp;
 378
 379         list_for_each_entry_safe(va_block, tmp, va_list, node) {
 380                 list_del(&va_block->node);
 381                 kfree(va_block);
 382         }
 383 }
 384
 385 /**
 386  * print_va_list_locked() - print virtual addresses list.
 387  * @hdev: habanalabs device structure.
 388  * @va_list: list of virtual addresses to print.
 389  *
 390  * This function does the following:
 391  * - Iterate over the list and print each virtual addresses block.
 392  *
 393  * This function should be called only when va_list lock is taken.
 394  */
 395 static void print_va_list_locked(struct hl_device *hdev,
 396                 struct list_head *va_list)
 397 {
 398 #if HL_MMU_DEBUG
 399         struct hl_vm_va_block *va_block;
 400
 401         dev_dbg(hdev->dev, "print va list:\n");
 402
 403         list_for_each_entry(va_block, va_list, node)
 404                 dev_dbg(hdev->dev,
 405                         "va block, start: 0x%llx, end: 0x%llx, size: %llu\n",
 406                         va_block->start, va_block->end, va_block->size);
 407 #endif
 408 }
 409
 410 /**
 411  * merge_va_blocks_locked() - merge a virtual block if possible.
 412  * @hdev: pointer to the habanalabs device structure.
 413  * @va_list: pointer to the virtual addresses block list.
 414  * @va_block: virtual block to merge with adjacent blocks.
 415  *
 416  * This function does the following:
 417  * - Merge the given blocks with the adjacent blocks if their virtual ranges
 418  *   create a contiguous virtual range.
 419  *
 420  * This Function should be called only when va_list lock is taken.
 421  */
 422 static void merge_va_blocks_locked(struct hl_device *hdev,
 423                 struct list_head *va_list, struct hl_vm_va_block *va_block)
 424 {
 425         struct hl_vm_va_block *prev, *next;
 426
 427         prev = list_prev_entry(va_block, node);
 428         if (&prev->node != va_list && prev->end + 1 == va_block->start) {
 429                 prev->end = va_block->end;
 430                 prev->size = prev->end - prev->start;
 431                 list_del(&va_block->node);
 432                 kfree(va_block);
 433                 va_block = prev;
 434         }
 435
 436         next = list_next_entry(va_block, node);
 437         if (&next->node != va_list && va_block->end + 1 == next->start) {
 438                 next->start = va_block->start;
 439                 next->size = next->end - next->start;
 440                 list_del(&va_block->node);
 441                 kfree(va_block);
 442         }
 443 }
 444
 445 /**
 446  * add_va_block_locked() - add a virtual block to the virtual addresses list.
 447  * @hdev: pointer to the habanalabs device structure.
 448  * @va_list: pointer to the virtual addresses block list.
 449  * @start: start virtual address.
 450  * @end: end virtual address.
 451  *
 452  * This function does the following:
 453  * - Add the given block to the virtual blocks list and merge with other blocks
 454  *   if a contiguous virtual block can be created.
 455  *
 456  * This Function should be called only when va_list lock is taken.
 457  */
 458 static int add_va_block_locked(struct hl_device *hdev,
 459                 struct list_head *va_list, u64 start, u64 end)
 460 {
 461         struct hl_vm_va_block *va_block, *res = NULL;
 462         u64 size = end - start;
 463
 464         print_va_list_locked(hdev, va_list);
 465
 466         list_for_each_entry(va_block, va_list, node) {
 467                 /* TODO: remove upon matureness */
 468                 if (hl_mem_area_crosses_range(start, size, va_block->start,
 469                                 va_block->end)) {
 470                         dev_err(hdev->dev,
 471                                 "block crossing ranges at start 0x%llx, end 0x%llx\n",
 472                                 va_block->start, va_block->end);
 473                         return -EINVAL;
 474                 }
 475
 476                 if (va_block->end < start)
 477                         res = va_block;
 478         }
 479
 480         va_block = kmalloc(sizeof(*va_block), GFP_KERNEL);
 481         if (!va_block)
 482                 return -ENOMEM;
 483
 484         va_block->start = start;
 485         va_block->end = end;
 486         va_block->size = size;
 487
 488         if (!res)
 489                 list_add(&va_block->node, va_list);
 490         else
 491                 list_add(&va_block->node, &res->node);
 492
 493         merge_va_blocks_locked(hdev, va_list, va_block);
 494
 495         print_va_list_locked(hdev, va_list);
 496
 497         return 0;
 498 }
 499
 500 /**
 501  * add_va_block() - wrapper for add_va_block_locked.
 502  * @hdev: pointer to the habanalabs device structure.
 503  * @va_list: pointer to the virtual addresses block list.
 504  * @start: start virtual address.
 505  * @end: end virtual address.
 506  *
 507  * This function does the following:
 508  * - Takes the list lock and calls add_va_block_locked.
 509  */
 510 static inline int add_va_block(struct hl_device *hdev,
 511                 struct hl_va_range *va_range, u64 start, u64 end)
 512 {
 513         int rc;
 514
 515         mutex_lock(&va_range->lock);
 516         rc = add_va_block_locked(hdev, &va_range->list, start, end);
 517         mutex_unlock(&va_range->lock);
 518
 519         return rc;
 520 }
 521
 522 /**
 523  * get_va_block() - get a virtual block for the given size and alignment.
 524  *
 525  * @hdev: pointer to the habanalabs device structure.
 526  * @va_range: pointer to the virtual addresses range.
 527  * @size: requested block size.
 528  * @hint_addr: hint for requested address by the user.
 529  * @va_block_align: required alignment of the virtual block start address.
 530  *
 531  * This function does the following:
 532  * - Iterate on the virtual block list to find a suitable virtual block for the
 533  *   given size, hint address and alignment.
 534  * - Reserve the requested block and update the list.
 535  * - Return the start address of the virtual block.
 536  */
 537 static u64 get_va_block(struct hl_device *hdev,
 538                                 struct hl_va_range *va_range,
 539                                 u64 size, u64 hint_addr, u32 va_block_align)
 540 {
 541         struct hl_vm_va_block *va_block, *new_va_block = NULL;
 542         u64 tmp_hint_addr, valid_start, valid_size, prev_start, prev_end,
 543                 align_mask, reserved_valid_start = 0, reserved_valid_size = 0;
 544         bool add_prev = false;
 545         bool is_align_pow_2  = is_power_of_2(va_range->page_size);
 546
 547         if (is_align_pow_2)
 548                 align_mask = ~((u64)va_block_align - 1);
 549         else
 550                 /*
 551                  * with non-power-of-2 range we work only with page granularity
 552                  * and the start address is page aligned,
 553                  * so no need for alignment checking.
 554                  */
 555                 size = DIV_ROUND_UP_ULL(size, va_range->page_size) *
 556                                                         va_range->page_size;
 557
 558         tmp_hint_addr = hint_addr;
 559
 560         /* Check if we need to ignore hint address */
 561         if ((is_align_pow_2 && (hint_addr & (va_block_align - 1))) ||
 562                         (!is_align_pow_2 &&
 563                                 do_div(tmp_hint_addr, va_range->page_size))) {
 564                 dev_info(hdev->dev, "Hint address 0x%llx will be ignored\n",
 565                                         hint_addr);
 566                 hint_addr = 0;
 567         }
 568
 569         mutex_lock(&va_range->lock);
 570
 571         print_va_list_locked(hdev, &va_range->list);
 572
 573         list_for_each_entry(va_block, &va_range->list, node) {
 574                 /* Calc the first possible aligned addr */
 575                 valid_start = va_block->start;
 576
 577                 if (is_align_pow_2 && (valid_start & (va_block_align - 1))) {
 578                         valid_start &= align_mask;
 579                         valid_start += va_block_align;
 580                         if (valid_start > va_block->end)
 581                                 continue;
 582                 }
 583
 584                 valid_size = va_block->end - valid_start;
 585                 if (valid_size < size)
 586                         continue;
 587
 588                 /* Pick the minimal length block which has the required size */
 589                 if (!new_va_block || (valid_size < reserved_valid_size)) {
 590                         new_va_block = va_block;
 591                         reserved_valid_start = valid_start;
 592                         reserved_valid_size = valid_size;
 593                 }
 594
 595                 if (hint_addr && hint_addr >= valid_start &&
 596                                         (hint_addr + size) <= va_block->end) {
 597                         new_va_block = va_block;
 598                         reserved_valid_start = hint_addr;
 599                         reserved_valid_size = valid_size;
 600                         break;
 601                 }
 602         }
 603
 604         if (!new_va_block) {
 605                 dev_err(hdev->dev, "no available va block for size %llu\n",
 606                                                                 size);
 607                 goto out;
 608         }
 609
 610         /*
 611          * Check if there is some leftover range due to reserving the new
 612          * va block, then return it to the main virtual addresses list.
 613          */
 614         if (reserved_valid_start > new_va_block->start) {
 615                 prev_start = new_va_block->start;
 616                 prev_end = reserved_valid_start - 1;
 617
 618                 new_va_block->start = reserved_valid_start;
 619                 new_va_block->size = reserved_valid_size;
 620
 621                 add_prev = true;
 622         }
 623
 624         if (new_va_block->size > size) {
 625                 new_va_block->start += size;
 626                 new_va_block->size = new_va_block->end - new_va_block->start;
 627         } else {
 628                 list_del(&new_va_block->node);
 629                 kfree(new_va_block);
 630         }
 631
 632         if (add_prev)
 633                 add_va_block_locked(hdev, &va_range->list, prev_start,
 634                                 prev_end);
 635
 636         print_va_list_locked(hdev, &va_range->list);
 637 out:
 638         mutex_unlock(&va_range->lock);
 639
 640         return reserved_valid_start;
 641 }
 642
 643 /*
 644  * hl_reserve_va_block() - reserve a virtual block of a given size.
 645  * @hdev: pointer to the habanalabs device structure.
 646  * @ctx: current context
 647  * @type: virtual addresses range type.
 648  * @size: requested block size.
 649  * @alignment: required alignment in bytes of the virtual block start address,
 650  *             0 means no alignment.
 651  *
 652  * This function does the following:
 653  * - Iterate on the virtual block list to find a suitable virtual block for the
 654  *   given size and alignment.
 655  * - Reserve the requested block and update the list.
 656  * - Return the start address of the virtual block.
 657  */
 658 u64 hl_reserve_va_block(struct hl_device *hdev, struct hl_ctx *ctx,
 659                 enum hl_va_range_type type, u32 size, u32 alignment)
 660 {
 661         return get_va_block(hdev, ctx->va_range[type], size, 0,
 662                         max(alignment, ctx->va_range[type]->page_size));
 663 }
 664
 665 /**
 666  * hl_get_va_range_type() - get va_range type for the given address and size.
 667  * @address: the start address of the area we want to validate.
 668  * @size: the size in bytes of the area we want to validate.
 669  * @type: returned va_range type.
 670  *
 671  * Return: true if the area is inside a valid range, false otherwise.
 672  */
 673 static int hl_get_va_range_type(struct hl_ctx *ctx, u64 address, u64 size,
 674                         enum hl_va_range_type *type)
 675 {
 676         int i;
 677
 678         for (i = 0 ; i < HL_VA_RANGE_TYPE_MAX; i++) {
 679                 if (hl_mem_area_inside_range(address, size,
 680                                 ctx->va_range[i]->start_addr,
 681                                 ctx->va_range[i]->end_addr)) {
 682                         *type = i;
 683                         return 0;
 684                 }
 685         }
 686
 687         return -EINVAL;
 688 }
 689
 690 /**
 691  * hl_unreserve_va_block() - wrapper for add_va_block to unreserve a va block.
 692  * @hdev: pointer to the habanalabs device structure
 693  * @ctx: pointer to the context structure.
 694  * @start: start virtual address.
 695  * @end: end virtual address.
 696  *
 697  * This function does the following:
 698  * - Takes the list lock and calls add_va_block_locked.
 699  */
 700 int hl_unreserve_va_block(struct hl_device *hdev, struct hl_ctx *ctx,
 701                 u64 start_addr, u64 size)
 702 {
 703         enum hl_va_range_type type;
 704         int rc;
 705
 706         rc = hl_get_va_range_type(ctx, start_addr, size, &type);
 707         if (rc) {
 708                 dev_err(hdev->dev,
 709                         "cannot find va_range for va %#llx size %llu",
 710                         start_addr, size);
 711                 return rc;
 712         }
 713
 714         rc = add_va_block(hdev, ctx->va_range[type], start_addr,
 715                                                 start_addr + size - 1);
 716         if (rc)
 717                 dev_warn(hdev->dev,
 718                         "add va block failed for vaddr: 0x%llx\n", start_addr);
 719
 720         return rc;
 721 }
 722
 723 /**
 724  * get_sg_info() - get number of pages and the DMA address from SG list.
 725  * @sg: the SG list.
 726  * @dma_addr: pointer to DMA address to return.
 727  *
 728  * Calculate the number of consecutive pages described by the SG list. Take the
 729  * offset of the address in the first page, add to it the length and round it up
 730  * to the number of needed pages.
 731  */
 732 static u32 get_sg_info(struct scatterlist *sg, dma_addr_t *dma_addr)
 733 {
 734         *dma_addr = sg_dma_address(sg);
 735
 736         return ((((*dma_addr) & (PAGE_SIZE - 1)) + sg_dma_len(sg)) +
 737                         (PAGE_SIZE - 1)) >> PAGE_SHIFT;
 738 }
 739
 740 /**
 741  * init_phys_pg_pack_from_userptr() - initialize physical page pack from host
 742  *                                    memory
 743  * @ctx: pointer to the context structure.
 744  * @userptr: userptr to initialize from.
 745  * @pphys_pg_pack: result pointer.
 746  *
 747  * This function does the following:
 748  * - Pin the physical pages related to the given virtual block.
 749  * - Create a physical page pack from the physical pages related to the given
 750  *   virtual block.
 751  */
 752 static int init_phys_pg_pack_from_userptr(struct hl_ctx *ctx,
 753                                 struct hl_userptr *userptr,
 754                                 struct hl_vm_phys_pg_pack **pphys_pg_pack)
 755 {
 756         struct hl_vm_phys_pg_pack *phys_pg_pack;
 757         struct scatterlist *sg;
 758         dma_addr_t dma_addr;
 759         u64 page_mask, total_npages;
 760         u32 npages, page_size = PAGE_SIZE,
 761                 huge_page_size = ctx->hdev->asic_prop.pmmu_huge.page_size;
 762         bool first = true, is_huge_page_opt = true;
 763         int rc, i, j;
 764         u32 pgs_in_huge_page = huge_page_size >> __ffs(page_size);
 765
 766         phys_pg_pack = kzalloc(sizeof(*phys_pg_pack), GFP_KERNEL);
 767         if (!phys_pg_pack)
 768                 return -ENOMEM;
 769
 770         phys_pg_pack->vm_type = userptr->vm_type;
 771         phys_pg_pack->created_from_userptr = true;
 772         phys_pg_pack->asid = ctx->asid;
 773         atomic_set(&phys_pg_pack->mapping_cnt, 1);
 774
 775         /* Only if all dma_addrs are aligned to 2MB and their
 776          * sizes is at least 2MB, we can use huge page mapping.
 777          * We limit the 2MB optimization to this condition,
 778          * since later on we acquire the related VA range as one
 779          * consecutive block.
 780          */
 781         total_npages = 0;
 782         for_each_sg(userptr->sgt->sgl, sg, userptr->sgt->nents, i) {
 783                 npages = get_sg_info(sg, &dma_addr);
 784
 785                 total_npages += npages;
 786
 787                 if ((npages % pgs_in_huge_page) ||
 788                                         (dma_addr & (huge_page_size - 1)))
 789                         is_huge_page_opt = false;
 790         }
 791
 792         if (is_huge_page_opt) {
 793                 page_size = huge_page_size;
 794                 do_div(total_npages, pgs_in_huge_page);
 795         }
 796
 797         page_mask = ~(((u64) page_size) - 1);
 798
 799         phys_pg_pack->pages = kvmalloc_array(total_npages, sizeof(u64),
 800                                                 GFP_KERNEL);
 801         if (ZERO_OR_NULL_PTR(phys_pg_pack->pages)) {
 802                 rc = -ENOMEM;
 803                 goto page_pack_arr_mem_err;
 804         }
 805
 806         phys_pg_pack->npages = total_npages;
 807         phys_pg_pack->page_size = page_size;
 808         phys_pg_pack->total_size = total_npages * page_size;
 809
 810         j = 0;
 811         for_each_sg(userptr->sgt->sgl, sg, userptr->sgt->nents, i) {
 812                 npages = get_sg_info(sg, &dma_addr);
 813
 814                 /* align down to physical page size and save the offset */
 815                 if (first) {
 816                         first = false;
 817                         phys_pg_pack->offset = dma_addr & (page_size - 1);
 818                         dma_addr &= page_mask;
 819                 }
 820
 821                 while (npages) {
 822                         phys_pg_pack->pages[j++] = dma_addr;
 823                         dma_addr += page_size;
 824
 825                         if (is_huge_page_opt)
 826                                 npages -= pgs_in_huge_page;
 827                         else
 828                                 npages--;
 829                 }
 830         }
 831
 832         *pphys_pg_pack = phys_pg_pack;
 833
 834         return 0;
 835
 836 page_pack_arr_mem_err:
 837         kfree(phys_pg_pack);
 838
 839         return rc;
 840 }
 841
 842 /**
 843  * map_phys_pg_pack() - maps the physical page pack..
 844  * @ctx: pointer to the context structure.
 845  * @vaddr: start address of the virtual area to map from.
 846  * @phys_pg_pack: the pack of physical pages to map to.
 847  *
 848  * This function does the following:
 849  * - Maps each chunk of virtual memory to matching physical chunk.
 850  * - Stores number of successful mappings in the given argument.
 851  * - Returns 0 on success, error code otherwise.
 852  */
 853 static int map_phys_pg_pack(struct hl_ctx *ctx, u64 vaddr,
 854                                 struct hl_vm_phys_pg_pack *phys_pg_pack)
 855 {
 856         struct hl_device *hdev = ctx->hdev;
 857         u64 next_vaddr = vaddr, paddr, mapped_pg_cnt = 0, i;
 858         u32 page_size = phys_pg_pack->page_size;
 859         int rc = 0;
 860
 861         for (i = 0 ; i < phys_pg_pack->npages ; i++) {
 862                 paddr = phys_pg_pack->pages[i];
 863
 864                 rc = hl_mmu_map_page(ctx, next_vaddr, paddr, page_size,
 865                                 (i + 1) == phys_pg_pack->npages);
 866                 if (rc) {
 867                         dev_err(hdev->dev,
 868                                 "map failed for handle %u, npages: %llu, mapped: %llu",
 869                                 phys_pg_pack->handle, phys_pg_pack->npages,
 870                                 mapped_pg_cnt);
 871                         goto err;
 872                 }
 873
 874                 mapped_pg_cnt++;
 875                 next_vaddr += page_size;
 876         }
 877
 878         return 0;
 879
 880 err:
 881         next_vaddr = vaddr;
 882         for (i = 0 ; i < mapped_pg_cnt ; i++) {
 883                 if (hl_mmu_unmap_page(ctx, next_vaddr, page_size,
 884                                         (i + 1) == mapped_pg_cnt))
 885                         dev_warn_ratelimited(hdev->dev,
 886                                 "failed to unmap handle %u, va: 0x%llx, pa: 0x%llx, page size: %u\n",
 887                                         phys_pg_pack->handle, next_vaddr,
 888                                         phys_pg_pack->pages[i], page_size);
 889
 890                 next_vaddr += page_size;
 891         }
 892
 893         return rc;
 894 }
 895
 896 /**
 897  * unmap_phys_pg_pack() - unmaps the physical page pack.
 898  * @ctx: pointer to the context structure.
 899  * @vaddr: start address of the virtual area to unmap.
 900  * @phys_pg_pack: the pack of physical pages to unmap.
 901  */
 902 static void unmap_phys_pg_pack(struct hl_ctx *ctx, u64 vaddr,
 903                                 struct hl_vm_phys_pg_pack *phys_pg_pack)
 904 {
 905         struct hl_device *hdev = ctx->hdev;
 906         u64 next_vaddr, i;
 907         bool is_host_addr;
 908         u32 page_size;
 909
 910         is_host_addr = !hl_is_dram_va(hdev, vaddr);
 911         page_size = phys_pg_pack->page_size;
 912         next_vaddr = vaddr;
 913
 914         for (i = 0 ; i < phys_pg_pack->npages ; i++, next_vaddr += page_size) {
 915                 if (hl_mmu_unmap_page(ctx, next_vaddr, page_size,
 916                                        (i + 1) == phys_pg_pack->npages))
 917                         dev_warn_ratelimited(hdev->dev,
 918                         "unmap failed for vaddr: 0x%llx\n", next_vaddr);
 919
 920                 /*
 921                  * unmapping on Palladium can be really long, so avoid a CPU
 922                  * soft lockup bug by sleeping a little between unmapping pages
 923                  *
 924                  * In addition, when unmapping host memory we pass through
 925                  * the Linux kernel to unpin the pages and that takes a long
 926                  * time. Therefore, sleep every 32K pages to avoid soft lockup
 927                  */
 928                 if (hdev->pldm || (is_host_addr && (i & 0x7FFF) == 0))
 929                         usleep_range(50, 200);
 930         }
 931 }
 932
 933 static int get_paddr_from_handle(struct hl_ctx *ctx, struct hl_mem_in *args,
 934                                         u64 *paddr)
 935 {
 936         struct hl_device *hdev = ctx->hdev;
 937         struct hl_vm *vm = &hdev->vm;
 938         struct hl_vm_phys_pg_pack *phys_pg_pack;
 939         u32 handle;
 940
 941         handle = lower_32_bits(args->map_device.handle);
 942         spin_lock(&vm->idr_lock);
 943         phys_pg_pack = idr_find(&vm->phys_pg_pack_handles, handle);
 944         if (!phys_pg_pack) {
 945                 spin_unlock(&vm->idr_lock);
 946                 dev_err(hdev->dev, "no match for handle %u\n", handle);
 947                 return -EINVAL;
 948         }
 949
 950         *paddr = phys_pg_pack->pages[0];
 951
 952         spin_unlock(&vm->idr_lock);
 953
 954         return 0;
 955 }
 956
 957 /**
 958  * map_device_va() - map the given memory.
 959  * @ctx: pointer to the context structure.
 960  * @args: host parameters with handle/host virtual address.
 961  * @device_addr: pointer to result device virtual address.
 962  *
 963  * This function does the following:
 964  * - If given a physical device memory handle, map to a device virtual block
 965  *   and return the start address of this block.
 966  * - If given a host virtual address and size, find the related physical pages,
 967  *   map a device virtual block to this pages and return the start address of
 968  *   this block.
 969  */
 970 static int map_device_va(struct hl_ctx *ctx, struct hl_mem_in *args,
 971                 u64 *device_addr)
 972 {
 973         struct hl_device *hdev = ctx->hdev;
 974         struct hl_vm *vm = &hdev->vm;
 975         struct hl_vm_phys_pg_pack *phys_pg_pack;
 976         struct hl_userptr *userptr = NULL;
 977         struct hl_vm_hash_node *hnode;
 978         struct hl_va_range *va_range;
 979         enum vm_type_t *vm_type;
 980         u64 ret_vaddr, hint_addr;
 981         u32 handle = 0, va_block_align;
 982         int rc;
 983         bool is_userptr = args->flags & HL_MEM_USERPTR;
 984
 985         /* Assume failure */
 986         *device_addr = 0;
 987
 988         if (is_userptr) {
 989                 u64 addr = args->map_host.host_virt_addr,
 990                         size = args->map_host.mem_size;
 991                 u32 page_size = hdev->asic_prop.pmmu.page_size,
 992                         huge_page_size = hdev->asic_prop.pmmu_huge.page_size;
 993
 994                 rc = dma_map_host_va(hdev, addr, size, &userptr);
 995                 if (rc) {
 996                         dev_err(hdev->dev, "failed to get userptr from va\n");
 997                         return rc;
 998                 }
 999
1000                 rc = init_phys_pg_pack_from_userptr(ctx, userptr,
1001                                 &phys_pg_pack);
1002                 if (rc) {
1003                         dev_err(hdev->dev,
1004                                 "unable to init page pack for vaddr 0x%llx\n",
1005                                 addr);
1006                         goto init_page_pack_err;
1007                 }
1008
1009                 vm_type = (enum vm_type_t *) userptr;
1010                 hint_addr = args->map_host.hint_addr;
1011                 handle = phys_pg_pack->handle;
1012
1013                 /* get required alignment */
1014                 if (phys_pg_pack->page_size == page_size) {
1015                         va_range = ctx->va_range[HL_VA_RANGE_TYPE_HOST];
1016
1017                         /*
1018                          * huge page alignment may be needed in case of regular
1019                          * page mapping, depending on the host VA alignment
1020                          */
1021                         if (addr & (huge_page_size - 1))
1022                                 va_block_align = page_size;
1023                         else
1024                                 va_block_align = huge_page_size;
1025                 } else {
1026                         /*
1027                          * huge page alignment is needed in case of huge page
1028                          * mapping
1029                          */
1030                         va_range = ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE];
1031                         va_block_align = huge_page_size;
1032                 }
1033         } else {
1034                 handle = lower_32_bits(args->map_device.handle);
1035
1036                 spin_lock(&vm->idr_lock);
1037                 phys_pg_pack = idr_find(&vm->phys_pg_pack_handles, handle);
1038                 if (!phys_pg_pack) {
1039                         spin_unlock(&vm->idr_lock);
1040                         dev_err(hdev->dev,
1041                                 "no match for handle %u\n", handle);
1042                         return -EINVAL;
1043                 }
1044
1045                 /* increment now to avoid freeing device memory while mapping */
1046                 atomic_inc(&phys_pg_pack->mapping_cnt);
1047
1048                 spin_unlock(&vm->idr_lock);
1049
1050                 vm_type = (enum vm_type_t *) phys_pg_pack;
1051
1052                 hint_addr = args->map_device.hint_addr;
1053
1054                 /* DRAM VA alignment is the same as the MMU page size */
1055                 va_range = ctx->va_range[HL_VA_RANGE_TYPE_DRAM];
1056                 va_block_align = hdev->asic_prop.dmmu.page_size;
1057         }
1058
1059         /*
1060          * relevant for mapping device physical memory only, as host memory is
1061          * implicitly shared
1062          */
1063         if (!is_userptr && !(phys_pg_pack->flags & HL_MEM_SHARED) &&
1064                         phys_pg_pack->asid != ctx->asid) {
1065                 dev_err(hdev->dev,
1066                         "Failed to map memory, handle %u is not shared\n",
1067                         handle);
1068                 rc = -EPERM;
1069                 goto shared_err;
1070         }
1071
1072         hnode = kzalloc(sizeof(*hnode), GFP_KERNEL);
1073         if (!hnode) {
1074                 rc = -ENOMEM;
1075                 goto hnode_err;
1076         }
1077
1078         ret_vaddr = get_va_block(hdev, va_range, phys_pg_pack->total_size,
1079                                         hint_addr, va_block_align);
1080         if (!ret_vaddr) {
1081                 dev_err(hdev->dev, "no available va block for handle %u\n",
1082                                 handle);
1083                 rc = -ENOMEM;
1084                 goto va_block_err;
1085         }
1086
1087         mutex_lock(&ctx->mmu_lock);
1088
1089         rc = map_phys_pg_pack(ctx, ret_vaddr, phys_pg_pack);
1090         if (rc) {
1091                 mutex_unlock(&ctx->mmu_lock);
1092                 dev_err(hdev->dev, "mapping page pack failed for handle %u\n",
1093                                 handle);
1094                 goto map_err;
1095         }
1096
1097         rc = hdev->asic_funcs->mmu_invalidate_cache(hdev, false, *vm_type);
1098
1099         mutex_unlock(&ctx->mmu_lock);
1100
1101         if (rc) {
1102                 dev_err(hdev->dev,
1103                         "mapping handle %u failed due to MMU cache invalidation\n",
1104                         handle);
1105                 goto map_err;
1106         }
1107
1108         ret_vaddr += phys_pg_pack->offset;
1109
1110         hnode->ptr = vm_type;
1111         hnode->vaddr = ret_vaddr;
1112
1113         mutex_lock(&ctx->mem_hash_lock);
1114         hash_add(ctx->mem_hash, &hnode->node, ret_vaddr);
1115         mutex_unlock(&ctx->mem_hash_lock);
1116
1117         *device_addr = ret_vaddr;
1118
1119         if (is_userptr)
1120                 free_phys_pg_pack(hdev, phys_pg_pack);
1121
1122         return 0;
1123
1124 map_err:
1125         if (add_va_block(hdev, va_range, ret_vaddr,
1126                                 ret_vaddr + phys_pg_pack->total_size - 1))
1127                 dev_warn(hdev->dev,
1128                         "release va block failed for handle 0x%x, vaddr: 0x%llx\n",
1129                                 handle, ret_vaddr);
1130
1131 va_block_err:
1132         kfree(hnode);
1133 hnode_err:
1134 shared_err:
1135         atomic_dec(&phys_pg_pack->mapping_cnt);
1136         if (is_userptr)
1137                 free_phys_pg_pack(hdev, phys_pg_pack);
1138 init_page_pack_err:
1139         if (is_userptr)
1140                 dma_unmap_host_va(hdev, userptr);
1141
1142         return rc;
1143 }
1144
1145 /**
1146  * unmap_device_va() - unmap the given device virtual address.
1147  * @ctx: pointer to the context structure.
1148  * @args: host parameters with device virtual address to unmap.
1149  * @ctx_free: true if in context free flow, false otherwise.
1150  *
1151  * This function does the following:
1152  * - unmap the physical pages related to the given virtual address.
1153  * - return the device virtual block to the virtual block list.
1154  */
1155 static int unmap_device_va(struct hl_ctx *ctx, struct hl_mem_in *args,
1156                                 bool ctx_free)
1157 {
1158         struct hl_device *hdev = ctx->hdev;
1159         struct asic_fixed_properties *prop = &hdev->asic_prop;
1160         struct hl_vm_phys_pg_pack *phys_pg_pack = NULL;
1161         struct hl_vm_hash_node *hnode = NULL;
1162         struct hl_userptr *userptr = NULL;
1163         struct hl_va_range *va_range;
1164         u64 vaddr = args->unmap.device_virt_addr;
1165         enum vm_type_t *vm_type;
1166         bool is_userptr;
1167         int rc = 0;
1168
1169         /* protect from double entrance */
1170         mutex_lock(&ctx->mem_hash_lock);
1171         hash_for_each_possible(ctx->mem_hash, hnode, node, (unsigned long)vaddr)
1172                 if (vaddr == hnode->vaddr)
1173                         break;
1174
1175         if (!hnode) {
1176                 mutex_unlock(&ctx->mem_hash_lock);
1177                 dev_err(hdev->dev,
1178                         "unmap failed, no mem hnode for vaddr 0x%llx\n",
1179                         vaddr);
1180                 return -EINVAL;
1181         }
1182
1183         hash_del(&hnode->node);
1184         mutex_unlock(&ctx->mem_hash_lock);
1185
1186         vm_type = hnode->ptr;
1187
1188         if (*vm_type == VM_TYPE_USERPTR) {
1189                 is_userptr = true;
1190                 userptr = hnode->ptr;
1191                 rc = init_phys_pg_pack_from_userptr(ctx, userptr,
1192                                                         &phys_pg_pack);
1193                 if (rc) {
1194                         dev_err(hdev->dev,
1195                                 "unable to init page pack for vaddr 0x%llx\n",
1196                                 vaddr);
1197                         goto vm_type_err;
1198                 }
1199
1200                 if (phys_pg_pack->page_size ==
1201                                         hdev->asic_prop.pmmu.page_size)
1202                         va_range = ctx->va_range[HL_VA_RANGE_TYPE_HOST];
1203                 else
1204                         va_range = ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE];
1205         } else if (*vm_type == VM_TYPE_PHYS_PACK) {
1206                 is_userptr = false;
1207                 va_range = ctx->va_range[HL_VA_RANGE_TYPE_DRAM];
1208                 phys_pg_pack = hnode->ptr;
1209         } else {
1210                 dev_warn(hdev->dev,
1211                         "unmap failed, unknown vm desc for vaddr 0x%llx\n",
1212                                 vaddr);
1213                 rc = -EFAULT;
1214                 goto vm_type_err;
1215         }
1216
1217         if (atomic_read(&phys_pg_pack->mapping_cnt) == 0) {
1218                 dev_err(hdev->dev, "vaddr 0x%llx is not mapped\n", vaddr);
1219                 rc = -EINVAL;
1220                 goto mapping_cnt_err;
1221         }
1222
1223         if (!is_userptr && !is_power_of_2(phys_pg_pack->page_size))
1224                 vaddr = prop->dram_base_address +
1225                         DIV_ROUND_DOWN_ULL(vaddr - prop->dram_base_address,
1226                                                 phys_pg_pack->page_size) *
1227                                                         phys_pg_pack->page_size;
1228         else
1229                 vaddr &= ~(((u64) phys_pg_pack->page_size) - 1);
1230
1231         mutex_lock(&ctx->mmu_lock);
1232
1233         unmap_phys_pg_pack(ctx, vaddr, phys_pg_pack);
1234
1235         /*
1236          * During context free this function is called in a loop to clean all
1237          * the context mappings. Hence the cache invalidation can be called once
1238          * at the loop end rather than for each iteration
1239          */
1240         if (!ctx_free)
1241                 rc = hdev->asic_funcs->mmu_invalidate_cache(hdev, true,
1242                                                                 *vm_type);
1243
1244         mutex_unlock(&ctx->mmu_lock);
1245
1246         /*
1247          * If the context is closing we don't need to check for the MMU cache
1248          * invalidation return code and update the VA free list as in this flow
1249          * we invalidate the MMU cache outside of this unmap function and the VA
1250          * free list will be freed anyway.
1251          */
1252         if (!ctx_free) {
1253                 int tmp_rc;
1254
1255                 if (rc)
1256                         dev_err(hdev->dev,
1257                                 "unmapping vaddr 0x%llx failed due to MMU cache invalidation\n",
1258                                 vaddr);
1259
1260                 tmp_rc = add_va_block(hdev, va_range, vaddr,
1261                                         vaddr + phys_pg_pack->total_size - 1);
1262                 if (tmp_rc) {
1263                         dev_warn(hdev->dev,
1264                                         "add va block failed for vaddr: 0x%llx\n",
1265                                         vaddr);
1266                         if (!rc)
1267                                 rc = tmp_rc;
1268                 }
1269         }
1270
1271         atomic_dec(&phys_pg_pack->mapping_cnt);
1272         kfree(hnode);
1273
1274         if (is_userptr) {
1275                 free_phys_pg_pack(hdev, phys_pg_pack);
1276                 dma_unmap_host_va(hdev, userptr);
1277         }
1278
1279         return rc;
1280
1281 mapping_cnt_err:
1282         if (is_userptr)
1283                 free_phys_pg_pack(hdev, phys_pg_pack);
1284 vm_type_err:
1285         mutex_lock(&ctx->mem_hash_lock);
1286         hash_add(ctx->mem_hash, &hnode->node, vaddr);
1287         mutex_unlock(&ctx->mem_hash_lock);
1288
1289         return rc;
1290 }
1291
1292 static int map_block(struct hl_device *hdev, u64 address, u64 *handle,
1293                         u32 *size)
1294 {
1295         u32 block_id = 0;
1296         int rc;
1297
1298         rc = hdev->asic_funcs->get_hw_block_id(hdev, address, size, &block_id);
1299
1300         *handle = block_id | HL_MMAP_TYPE_BLOCK;
1301         *handle <<= PAGE_SHIFT;
1302
1303         return rc;
1304 }
1305
1306 static void hw_block_vm_close(struct vm_area_struct *vma)
1307 {
1308         struct hl_ctx *ctx = (struct hl_ctx *) vma->vm_private_data;
1309
1310         hl_ctx_put(ctx);
1311         vma->vm_private_data = NULL;
1312 }
1313
1314 static const struct vm_operations_struct hw_block_vm_ops = {
1315         .close = hw_block_vm_close
1316 };
1317
1318 /**
1319  * hl_hw_block_mmap() - mmap a hw block to user.
1320  * @hpriv: pointer to the private data of the fd
1321  * @vma: pointer to vm_area_struct of the process
1322  *
1323  * Driver increments context reference for every HW block mapped in order
1324  * to prevent user from closing FD without unmapping first
1325  */
1326 int hl_hw_block_mmap(struct hl_fpriv *hpriv, struct vm_area_struct *vma)
1327 {
1328         struct hl_device *hdev = hpriv->hdev;
1329         u32 block_id, block_size;
1330         int rc;
1331
1332         /* We use the page offset to hold the block id and thus we need to clear
1333          * it before doing the mmap itself
1334          */
1335         block_id = vma->vm_pgoff;
1336         vma->vm_pgoff = 0;
1337
1338         /* Driver only allows mapping of a complete HW block */
1339         block_size = vma->vm_end - vma->vm_start;
1340
1341 #ifdef _HAS_TYPE_ARG_IN_ACCESS_OK
1342         if (!access_ok(VERIFY_WRITE,
1343                 (void __user *) (uintptr_t) vma->vm_start, block_size)) {
1344 #else
1345         if (!access_ok((void __user *) (uintptr_t) vma->vm_start, block_size)) {
1346 #endif
1347                 dev_err(hdev->dev,
1348                         "user pointer is invalid - 0x%lx\n",
1349                         vma->vm_start);
1350
1351                 return -EINVAL;
1352         }
1353
1354         vma->vm_ops = &hw_block_vm_ops;
1355         vma->vm_private_data = hpriv->ctx;
1356
1357         hl_ctx_get(hdev, hpriv->ctx);
1358
1359         rc = hdev->asic_funcs->hw_block_mmap(hdev, vma, block_id, block_size);
1360         if (rc) {
1361                 hl_ctx_put(hpriv->ctx);
1362                 return rc;
1363         }
1364
1365         vma->vm_pgoff = block_id;
1366
1367         return 0;
1368 }
1369
1370 static int mem_ioctl_no_mmu(struct hl_fpriv *hpriv, union hl_mem_args *args)
1371 {
1372         struct hl_device *hdev = hpriv->hdev;
1373         struct hl_ctx *ctx = hpriv->ctx;
1374         u64 block_handle, device_addr = 0;
1375         u32 handle = 0, block_size;
1376         int rc;
1377
1378         switch (args->in.op) {
1379         case HL_MEM_OP_ALLOC:
1380                 if (args->in.alloc.mem_size == 0) {
1381                         dev_err(hdev->dev,
1382                                 "alloc size must be larger than 0\n");
1383                         rc = -EINVAL;
1384                         goto out;
1385                 }
1386
1387                 /* Force contiguous as there are no real MMU
1388                  * translations to overcome physical memory gaps
1389                  */
1390                 args->in.flags |= HL_MEM_CONTIGUOUS;
1391                 rc = alloc_device_memory(ctx, &args->in, &handle);
1392
1393                 memset(args, 0, sizeof(*args));
1394                 args->out.handle = (__u64) handle;
1395                 break;
1396
1397         case HL_MEM_OP_FREE:
1398                 rc = free_device_memory(ctx, &args->in);
1399                 break;
1400
1401         case HL_MEM_OP_MAP:
1402                 if (args->in.flags & HL_MEM_USERPTR) {
1403                         device_addr = args->in.map_host.host_virt_addr;
1404                         rc = 0;
1405                 } else {
1406                         rc = get_paddr_from_handle(ctx, &args->in,
1407                                                         &device_addr);
1408                 }
1409
1410                 memset(args, 0, sizeof(*args));
1411                 args->out.device_virt_addr = device_addr;
1412                 break;
1413
1414         case HL_MEM_OP_UNMAP:
1415                 rc = 0;
1416                 break;
1417
1418         case HL_MEM_OP_MAP_BLOCK:
1419                 rc = map_block(hdev, args->in.map_block.block_addr,
1420                                 &block_handle, &block_size);
1421                 args->out.block_handle = block_handle;
1422                 args->out.block_size = block_size;
1423                 break;
1424
1425         default:
1426                 dev_err(hdev->dev, "Unknown opcode for memory IOCTL\n");
1427                 rc = -ENOTTY;
1428                 break;
1429         }
1430
1431 out:
1432         return rc;
1433 }
1434
1435 int hl_mem_ioctl(struct hl_fpriv *hpriv, void *data)
1436 {
1437         enum hl_device_status status;
1438         union hl_mem_args *args = data;
1439         struct hl_device *hdev = hpriv->hdev;
1440         struct hl_ctx *ctx = hpriv->ctx;
1441         u64 block_handle, device_addr = 0;
1442         u32 handle = 0, block_size;
1443         int rc;
1444
1445         if (!hl_device_operational(hdev, &status)) {
1446                 dev_warn_ratelimited(hdev->dev,
1447                         "Device is %s. Can't execute MEMORY IOCTL\n",
1448                         hdev->status[status]);
1449                 return -EBUSY;
1450         }
1451
1452         if (!hdev->mmu_enable)
1453                 return mem_ioctl_no_mmu(hpriv, args);
1454
1455         switch (args->in.op) {
1456         case HL_MEM_OP_ALLOC:
1457                 if (args->in.alloc.mem_size == 0) {
1458                         dev_err(hdev->dev,
1459                                 "alloc size must be larger than 0\n");
1460                         rc = -EINVAL;
1461                         goto out;
1462                 }
1463
1464                 /* If DRAM does not support virtual memory the driver won't
1465                  * handle the allocation/freeing of that memory. However, for
1466                  * system administration/monitoring purposes, the driver will
1467                  * keep track of the amount of DRAM memory that is allocated
1468                  * and freed by the user. Because this code totally relies on
1469                  * the user's input, the driver can't ensure the validity
1470                  * of this accounting.
1471                  */
1472                 if (!hdev->asic_prop.dram_supports_virtual_memory) {
1473                         atomic64_add(args->in.alloc.mem_size,
1474                                         &ctx->dram_phys_mem);
1475                         atomic64_add(args->in.alloc.mem_size,
1476                                         &hdev->dram_used_mem);
1477
1478                         dev_dbg(hdev->dev, "DRAM alloc is not supported\n");
1479                         rc = 0;
1480
1481                         memset(args, 0, sizeof(*args));
1482                         args->out.handle = 0;
1483                         goto out;
1484                 }
1485
1486                 rc = alloc_device_memory(ctx, &args->in, &handle);
1487
1488                 memset(args, 0, sizeof(*args));
1489                 args->out.handle = (__u64) handle;
1490                 break;
1491
1492         case HL_MEM_OP_FREE:
1493                 /* If DRAM does not support virtual memory the driver won't
1494                  * handle the allocation/freeing of that memory. However, for
1495                  * system administration/monitoring purposes, the driver will
1496                  * keep track of the amount of DRAM memory that is allocated
1497                  * and freed by the user. Because this code totally relies on
1498                  * the user's input, the driver can't ensure the validity
1499                  * of this accounting.
1500                  */
1501                 if (!hdev->asic_prop.dram_supports_virtual_memory) {
1502                         atomic64_sub(args->in.alloc.mem_size,
1503                                         &ctx->dram_phys_mem);
1504                         atomic64_sub(args->in.alloc.mem_size,
1505                                         &hdev->dram_used_mem);
1506
1507                         dev_dbg(hdev->dev, "DRAM alloc is not supported\n");
1508                         rc = 0;
1509
1510                         goto out;
1511                 }
1512
1513                 rc = free_device_memory(ctx, &args->in);
1514                 break;
1515
1516         case HL_MEM_OP_MAP:
1517                 rc = map_device_va(ctx, &args->in, &device_addr);
1518
1519                 memset(args, 0, sizeof(*args));
1520                 args->out.device_virt_addr = device_addr;
1521                 break;
1522
1523         case HL_MEM_OP_UNMAP:
1524                 rc = unmap_device_va(ctx, &args->in, false);
1525                 break;
1526
1527         case HL_MEM_OP_MAP_BLOCK:
1528                 rc = map_block(hdev, args->in.map_block.block_addr,
1529                                 &block_handle, &block_size);
1530                 args->out.block_handle = block_handle;
1531                 args->out.block_size = block_size;
1532                 break;
1533
1534         default:
1535                 dev_err(hdev->dev, "Unknown opcode for memory IOCTL\n");
1536                 rc = -ENOTTY;
1537                 break;
1538         }
1539
1540 out:
1541         return rc;
1542 }
1543
1544 static int get_user_memory(struct hl_device *hdev, u64 addr, u64 size,
1545                                 u32 npages, u64 start, u32 offset,
1546                                 struct hl_userptr *userptr)
1547 {
1548         int rc;
1549
1550         if (!access_ok((void __user *) (uintptr_t) addr, size)) {
1551                 dev_err(hdev->dev, "user pointer is invalid - 0x%llx\n", addr);
1552                 return -EFAULT;
1553         }
1554
1555         userptr->pages = kvmalloc_array(npages, sizeof(*userptr->pages),
1556                                         GFP_KERNEL);
1557         if (!userptr->pages)
1558                 return -ENOMEM;
1559
1560         rc = pin_user_pages_fast(start, npages,
1561                                  FOLL_FORCE | FOLL_WRITE | FOLL_LONGTERM,
1562                                  userptr->pages);
1563
1564         if (rc != npages) {
1565                 dev_err(hdev->dev,
1566                         "Failed to map host memory, user ptr probably wrong\n");
1567                 if (rc < 0)
1568                         goto destroy_pages;
1569                 npages = rc;
1570                 rc = -EFAULT;
1571                 goto put_pages;
1572         }
1573         userptr->npages = npages;
1574
1575         rc = sg_alloc_table_from_pages(userptr->sgt,
1576                                        userptr->pages,
1577                                        npages, offset, size, GFP_ATOMIC);
1578         if (rc < 0) {
1579                 dev_err(hdev->dev, "failed to create SG table from pages\n");
1580                 goto put_pages;
1581         }
1582
1583         return 0;
1584
1585 put_pages:
1586         unpin_user_pages(userptr->pages, npages);
1587 destroy_pages:
1588         kvfree(userptr->pages);
1589         return rc;
1590 }
1591
1592 /**
1593  * hl_pin_host_memory() - pins a chunk of host memory.
1594  * @hdev: pointer to the habanalabs device structure.
1595  * @addr: the host virtual address of the memory area.
1596  * @size: the size of the memory area.
1597  * @userptr: pointer to hl_userptr structure.
1598  *
1599  * This function does the following:
1600  * - Pins the physical pages.
1601  * - Create an SG list from those pages.
1602  */
1603 int hl_pin_host_memory(struct hl_device *hdev, u64 addr, u64 size,
1604                                         struct hl_userptr *userptr)
1605 {
1606         u64 start, end;
1607         u32 npages, offset;
1608         int rc;
1609
1610         if (!size) {
1611                 dev_err(hdev->dev, "size to pin is invalid - %llu\n", size);
1612                 return -EINVAL;
1613         }
1614
1615         /*
1616          * If the combination of the address and size requested for this memory
1617          * region causes an integer overflow, return error.
1618          */
1619         if (((addr + size) < addr) ||
1620                         PAGE_ALIGN(addr + size) < (addr + size)) {
1621                 dev_err(hdev->dev,
1622                         "user pointer 0x%llx + %llu causes integer overflow\n",
1623                         addr, size);
1624                 return -EINVAL;
1625         }
1626
1627         /*
1628          * This function can be called also from data path, hence use atomic
1629          * always as it is not a big allocation.
1630          */
1631         userptr->sgt = kzalloc(sizeof(*userptr->sgt), GFP_ATOMIC);
1632         if (!userptr->sgt)
1633                 return -ENOMEM;
1634
1635         start = addr & PAGE_MASK;
1636         offset = addr & ~PAGE_MASK;
1637         end = PAGE_ALIGN(addr + size);
1638         npages = (end - start) >> PAGE_SHIFT;
1639
1640         userptr->size = size;
1641         userptr->addr = addr;
1642         userptr->dma_mapped = false;
1643         INIT_LIST_HEAD(&userptr->job_node);
1644
1645         rc = get_user_memory(hdev, addr, size, npages, start, offset,
1646                                 userptr);
1647         if (rc) {
1648                 dev_err(hdev->dev,
1649                         "failed to get user memory for address 0x%llx\n",
1650                         addr);
1651                 goto free_sgt;
1652         }
1653
1654         hl_debugfs_add_userptr(hdev, userptr);
1655
1656         return 0;
1657
1658 free_sgt:
1659         kfree(userptr->sgt);
1660         return rc;
1661 }
1662
1663 /*
1664  * hl_unpin_host_memory - unpins a chunk of host memory.
1665  * @hdev: pointer to the habanalabs device structure
1666  * @userptr: pointer to hl_userptr structure
1667  *
1668  * This function does the following:
1669  * - Unpins the physical pages related to the host memory
1670  * - Free the SG list
1671  */
1672 void hl_unpin_host_memory(struct hl_device *hdev, struct hl_userptr *userptr)
1673 {
1674         hl_debugfs_remove_userptr(hdev, userptr);
1675
1676         if (userptr->dma_mapped)
1677                 hdev->asic_funcs->hl_dma_unmap_sg(hdev, userptr->sgt->sgl,
1678                                                         userptr->sgt->nents,
1679                                                         userptr->dir);
1680
1681         unpin_user_pages_dirty_lock(userptr->pages, userptr->npages, true);
1682         kvfree(userptr->pages);
1683
1684         list_del(&userptr->job_node);
1685
1686         sg_free_table(userptr->sgt);
1687         kfree(userptr->sgt);
1688 }
1689
1690 /**
1691  * hl_userptr_delete_list() - clear userptr list.
1692  * @hdev: pointer to the habanalabs device structure.
1693  * @userptr_list: pointer to the list to clear.
1694  *
1695  * This function does the following:
1696  * - Iterates over the list and unpins the host memory and frees the userptr
1697  *   structure.
1698  */
1699 void hl_userptr_delete_list(struct hl_device *hdev,
1700                                 struct list_head *userptr_list)
1701 {
1702         struct hl_userptr *userptr, *tmp;
1703
1704         list_for_each_entry_safe(userptr, tmp, userptr_list, job_node) {
1705                 hl_unpin_host_memory(hdev, userptr);
1706                 kfree(userptr);
1707         }
1708
1709         INIT_LIST_HEAD(userptr_list);
1710 }
1711
1712 /**
1713  * hl_userptr_is_pinned() - returns whether the given userptr is pinned.
1714  * @hdev: pointer to the habanalabs device structure.
1715  * @userptr_list: pointer to the list to clear.
1716  * @userptr: pointer to userptr to check.
1717  *
1718  * This function does the following:
1719  * - Iterates over the list and checks if the given userptr is in it, means is
1720  *   pinned. If so, returns true, otherwise returns false.
1721  */
1722 bool hl_userptr_is_pinned(struct hl_device *hdev, u64 addr,
1723                                 u32 size, struct list_head *userptr_list,
1724                                 struct hl_userptr **userptr)
1725 {
1726         list_for_each_entry((*userptr), userptr_list, job_node) {
1727                 if ((addr == (*userptr)->addr) && (size == (*userptr)->size))
1728                         return true;
1729         }
1730
1731         return false;
1732 }
1733
1734 /**
1735  * va_range_init() - initialize virtual addresses range.
1736  * @hdev: pointer to the habanalabs device structure.
1737  * @va_range: pointer to the range to initialize.
1738  * @start: range start address.
1739  * @end: range end address.
1740  *
1741  * This function does the following:
1742  * - Initializes the virtual addresses list of the given range with the given
1743  *   addresses.
1744  */
1745 static int va_range_init(struct hl_device *hdev, struct hl_va_range *va_range,
1746                                 u64 start, u64 end, u32 page_size)
1747 {
1748         int rc;
1749
1750         INIT_LIST_HEAD(&va_range->list);
1751
1752         /*
1753          * PAGE_SIZE alignment
1754          * it is the callers responsibility to align the addresses if the
1755          * page size is not a power of 2
1756          */
1757
1758         if (is_power_of_2(page_size)) {
1759                 if (start & (PAGE_SIZE - 1)) {
1760                         start &= PAGE_MASK;
1761                         start += PAGE_SIZE;
1762                 }
1763
1764                 if (end & (PAGE_SIZE - 1))
1765                         end &= PAGE_MASK;
1766         }
1767
1768         if (start >= end) {
1769                 dev_err(hdev->dev, "too small vm range for va list\n");
1770                 return -EFAULT;
1771         }
1772
1773         rc = add_va_block(hdev, va_range, start, end);
1774
1775         if (rc) {
1776                 dev_err(hdev->dev, "Failed to init host va list\n");
1777                 return rc;
1778         }
1779
1780         va_range->start_addr = start;
1781         va_range->end_addr = end;
1782         va_range->page_size = page_size;
1783
1784         return 0;
1785 }
1786
1787 /**
1788  * va_range_fini() - clear a virtual addresses range.
1789  * @hdev: pointer to the habanalabs structure.
1790  * va_range: pointer to virtual addresses rang.e
1791  *
1792  * This function does the following:
1793  * - Frees the virtual addresses block list and its lock.
1794  */
1795 static void va_range_fini(struct hl_device *hdev, struct hl_va_range *va_range)
1796 {
1797         mutex_lock(&va_range->lock);
1798         clear_va_list_locked(hdev, &va_range->list);
1799         mutex_unlock(&va_range->lock);
1800
1801         mutex_destroy(&va_range->lock);
1802         kfree(va_range);
1803 }
1804
1805 /**
1806  * vm_ctx_init_with_ranges() - initialize virtual memory for context.
1807  * @ctx: pointer to the habanalabs context structure.
1808  * @host_range_start: host virtual addresses range start.
1809  * @host_range_end: host virtual addresses range end.
1810  * @host_huge_range_start: host virtual addresses range start for memory
1811  *                         allocated with huge pages.
1812  * @host_huge_range_end: host virtual addresses range end for memory allocated
1813  *                        with huge pages.
1814  * @dram_range_start: dram virtual addresses range start.
1815  * @dram_range_end: dram virtual addresses range end.
1816  *
1817  * This function initializes the following:
1818  * - MMU for context.
1819  * - Virtual address to area descriptor hashtable.
1820  * - Virtual block list of available virtual memory.
1821  */
1822 static int vm_ctx_init_with_ranges(struct hl_ctx *ctx,
1823                                         u64 host_range_start,
1824                                         u64 host_range_end,
1825                                         u32 host_page_size,
1826                                         u64 host_huge_range_start,
1827                                         u64 host_huge_range_end,
1828                                         u32 host_huge_page_size,
1829                                         u64 dram_range_start,
1830                                         u64 dram_range_end,
1831                                         u32 dram_page_size)
1832 {
1833         struct hl_device *hdev = ctx->hdev;
1834         int i, rc;
1835
1836         for (i = 0 ; i < HL_VA_RANGE_TYPE_MAX ; i++) {
1837                 ctx->va_range[i] =
1838                         kzalloc(sizeof(struct hl_va_range), GFP_KERNEL);
1839                 if (!ctx->va_range[i]) {
1840                         rc = -ENOMEM;
1841                         goto free_va_range;
1842                 }
1843         }
1844
1845         rc = hl_mmu_ctx_init(ctx);
1846         if (rc) {
1847                 dev_err(hdev->dev, "failed to init context %d\n", ctx->asid);
1848                 goto free_va_range;
1849         }
1850
1851         mutex_init(&ctx->mem_hash_lock);
1852         hash_init(ctx->mem_hash);
1853
1854         mutex_init(&ctx->va_range[HL_VA_RANGE_TYPE_HOST]->lock);
1855
1856         rc = va_range_init(hdev, ctx->va_range[HL_VA_RANGE_TYPE_HOST],
1857                         host_range_start, host_range_end, host_page_size);
1858         if (rc) {
1859                 dev_err(hdev->dev, "failed to init host vm range\n");
1860                 goto mmu_ctx_fini;
1861         }
1862
1863         if (hdev->pmmu_huge_range) {
1864                 mutex_init(&ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]->lock);
1865
1866                 rc = va_range_init(hdev,
1867                         ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE],
1868                         host_huge_range_start, host_huge_range_end,
1869                         host_huge_page_size);
1870                 if (rc) {
1871                         dev_err(hdev->dev,
1872                                 "failed to init host huge vm range\n");
1873                         goto clear_host_va_range;
1874                 }
1875         } else {
1876                 kfree(ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]);
1877                 ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE] =
1878                                 ctx->va_range[HL_VA_RANGE_TYPE_HOST];
1879         }
1880
1881         mutex_init(&ctx->va_range[HL_VA_RANGE_TYPE_DRAM]->lock);
1882
1883         rc = va_range_init(hdev, ctx->va_range[HL_VA_RANGE_TYPE_DRAM],
1884                         dram_range_start, dram_range_end, dram_page_size);
1885         if (rc) {
1886                 dev_err(hdev->dev, "failed to init dram vm range\n");
1887                 goto clear_host_huge_va_range;
1888         }
1889
1890         hl_debugfs_add_ctx_mem_hash(hdev, ctx);
1891
1892         return 0;
1893
1894 clear_host_huge_va_range:
1895         mutex_destroy(&ctx->va_range[HL_VA_RANGE_TYPE_DRAM]->lock);
1896
1897         if (hdev->pmmu_huge_range) {
1898                 mutex_lock(&ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]->lock);
1899                 clear_va_list_locked(hdev,
1900                         &ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]->list);
1901                 mutex_unlock(&ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]->lock);
1902         }
1903 clear_host_va_range:
1904         if (hdev->pmmu_huge_range)
1905                 mutex_destroy(&ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]->lock);
1906         mutex_lock(&ctx->va_range[HL_VA_RANGE_TYPE_HOST]->lock);
1907         clear_va_list_locked(hdev, &ctx->va_range[HL_VA_RANGE_TYPE_HOST]->list);
1908         mutex_unlock(&ctx->va_range[HL_VA_RANGE_TYPE_HOST]->lock);
1909 mmu_ctx_fini:
1910         mutex_destroy(&ctx->va_range[HL_VA_RANGE_TYPE_HOST]->lock);
1911         mutex_destroy(&ctx->mem_hash_lock);
1912         hl_mmu_ctx_fini(ctx);
1913 free_va_range:
1914         for (i = 0 ; i < HL_VA_RANGE_TYPE_MAX ; i++)
1915                 kfree(ctx->va_range[i]);
1916
1917         return rc;
1918 }
1919
1920 int hl_vm_ctx_init(struct hl_ctx *ctx)
1921 {
1922         struct asic_fixed_properties *prop = &ctx->hdev->asic_prop;
1923         u64 host_range_start, host_range_end, host_huge_range_start,
1924                 host_huge_range_end, dram_range_start, dram_range_end;
1925         u32 host_page_size, host_huge_page_size, dram_page_size;
1926
1927         atomic64_set(&ctx->dram_phys_mem, 0);
1928
1929         /*
1930          * - If MMU is enabled, init the ranges as usual.
1931          * - If MMU is disabled, in case of host mapping, the returned address
1932          *   is the given one.
1933          *   In case of DRAM mapping, the returned address is the physical
1934          *   address of the memory related to the given handle.
1935          */
1936         if (!ctx->hdev->mmu_enable)
1937                 return 0;
1938
1939         dram_range_start = prop->dmmu.start_addr;
1940         dram_range_end = prop->dmmu.end_addr;
1941         dram_page_size = prop->dram_page_size ?
1942                                 prop->dram_page_size : prop->dmmu.page_size;
1943         host_range_start = prop->pmmu.start_addr;
1944         host_range_end = prop->pmmu.end_addr;
1945         host_page_size = prop->pmmu.page_size;
1946         host_huge_range_start = prop->pmmu_huge.start_addr;
1947         host_huge_range_end = prop->pmmu_huge.end_addr;
1948         host_huge_page_size = prop->pmmu_huge.page_size;
1949
1950         return vm_ctx_init_with_ranges(ctx, host_range_start, host_range_end,
1951                         host_page_size, host_huge_range_start,
1952                         host_huge_range_end, host_huge_page_size,
1953                         dram_range_start, dram_range_end, dram_page_size);
1954 }
1955
1956 /**
1957  * hl_vm_ctx_fini() - virtual memory teardown of context.
1958  * @ctx: pointer to the habanalabs context structure.
1959  *
1960  * This function perform teardown the following:
1961  * - Virtual block list of available virtual memory.
1962  * - Virtual address to area descriptor hashtable.
1963  * - MMU for context.
1964  *
1965  * In addition this function does the following:
1966  * - Unmaps the existing hashtable nodes if the hashtable is not empty. The
1967  *   hashtable should be empty as no valid mappings should exist at this
1968  *   point.
1969  * - Frees any existing physical page list from the idr which relates to the
1970  *   current context asid.
1971  * - This function checks the virtual block list for correctness. At this point
1972  *   the list should contain one element which describes the whole virtual
1973  *   memory range of the context. Otherwise, a warning is printed.
1974  */
1975 void hl_vm_ctx_fini(struct hl_ctx *ctx)
1976 {
1977         struct hl_device *hdev = ctx->hdev;
1978         struct hl_vm *vm = &hdev->vm;
1979         struct hl_vm_phys_pg_pack *phys_pg_list;
1980         struct hl_vm_hash_node *hnode;
1981         struct hlist_node *tmp_node;
1982         struct hl_mem_in args;
1983         int i;
1984
1985         if (!hdev->mmu_enable)
1986                 return;
1987
1988         hl_debugfs_remove_ctx_mem_hash(hdev, ctx);
1989
1990         /*
1991          * Clearly something went wrong on hard reset so no point in printing
1992          * another side effect error
1993          */
1994         if (!hdev->hard_reset_pending && !hash_empty(ctx->mem_hash))
1995                 dev_notice(hdev->dev,
1996                         "user released device without removing its memory mappings\n");
1997
1998         hash_for_each_safe(ctx->mem_hash, i, tmp_node, hnode, node) {
1999                 dev_dbg(hdev->dev,
2000                         "hl_mem_hash_node of vaddr 0x%llx of asid %d is still alive\n",
2001                         hnode->vaddr, ctx->asid);
2002                 args.unmap.device_virt_addr = hnode->vaddr;
2003                 unmap_device_va(ctx, &args, true);
2004         }
2005
2006         mutex_lock(&ctx->mmu_lock);
2007
2008         /* invalidate the cache once after the unmapping loop */
2009         hdev->asic_funcs->mmu_invalidate_cache(hdev, true, VM_TYPE_USERPTR);
2010         hdev->asic_funcs->mmu_invalidate_cache(hdev, true, VM_TYPE_PHYS_PACK);
2011
2012         mutex_unlock(&ctx->mmu_lock);
2013
2014         spin_lock(&vm->idr_lock);
2015         idr_for_each_entry(&vm->phys_pg_pack_handles, phys_pg_list, i)
2016                 if (phys_pg_list->asid == ctx->asid) {
2017                         dev_dbg(hdev->dev,
2018                                 "page list 0x%px of asid %d is still alive\n",
2019                                 phys_pg_list, ctx->asid);
2020                         atomic64_sub(phys_pg_list->total_size,
2021                                         &hdev->dram_used_mem);
2022                         free_phys_pg_pack(hdev, phys_pg_list);
2023                         idr_remove(&vm->phys_pg_pack_handles, i);
2024                 }
2025         spin_unlock(&vm->idr_lock);
2026
2027         va_range_fini(hdev, ctx->va_range[HL_VA_RANGE_TYPE_DRAM]);
2028         va_range_fini(hdev, ctx->va_range[HL_VA_RANGE_TYPE_HOST]);
2029
2030         if (hdev->pmmu_huge_range)
2031                 va_range_fini(hdev, ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]);
2032
2033         mutex_destroy(&ctx->mem_hash_lock);
2034         hl_mmu_ctx_fini(ctx);
2035
2036         /* In this case we need to clear the global accounting of DRAM usage
2037          * because the user notifies us on allocations. If the user is no more,
2038          * all DRAM is available
2039          */
2040         if (ctx->asid != HL_KERNEL_ASID_ID &&
2041                         !hdev->asic_prop.dram_supports_virtual_memory)
2042                 atomic64_set(&hdev->dram_used_mem, 0);
2043 }
2044
2045 /**
2046  * hl_vm_init() - initialize virtual memory module.
2047  * @hdev: pointer to the habanalabs device structure.
2048  *
2049  * This function initializes the following:
2050  * - MMU module.
2051  * - DRAM physical pages pool of 2MB.
2052  * - Idr for device memory allocation handles.
2053  */
2054 int hl_vm_init(struct hl_device *hdev)
2055 {
2056         struct asic_fixed_properties *prop = &hdev->asic_prop;
2057         struct hl_vm *vm = &hdev->vm;
2058         int rc;
2059
2060         if (is_power_of_2(prop->dram_page_size))
2061                 vm->dram_pg_pool =
2062                         gen_pool_create(__ffs(prop->dram_page_size), -1);
2063         else
2064                 vm->dram_pg_pool =
2065                         gen_pool_create(__ffs(DRAM_POOL_PAGE_SIZE), -1);
2066
2067         if (!vm->dram_pg_pool) {
2068                 dev_err(hdev->dev, "Failed to create dram page pool\n");
2069                 return -ENOMEM;
2070         }
2071
2072         kref_init(&vm->dram_pg_pool_refcount);
2073
2074         rc = gen_pool_add(vm->dram_pg_pool, prop->dram_user_base_address,
2075                         prop->dram_end_address - prop->dram_user_base_address,
2076                         -1);
2077
2078         if (rc) {
2079                 dev_err(hdev->dev,
2080                         "Failed to add memory to dram page pool %d\n", rc);
2081                 goto pool_add_err;
2082         }
2083
2084         spin_lock_init(&vm->idr_lock);
2085         idr_init(&vm->phys_pg_pack_handles);
2086
2087         atomic64_set(&hdev->dram_used_mem, 0);
2088
2089         vm->init_done = true;
2090
2091         return 0;
2092
2093 pool_add_err:
2094         gen_pool_destroy(vm->dram_pg_pool);
2095
2096         return rc;
2097 }
2098
2099 /**
2100  * hl_vm_fini() - virtual memory module teardown.
2101  * @hdev: pointer to the habanalabs device structure.
2102  *
2103  * This function perform teardown to the following:
2104  * - Idr for device memory allocation handles.
2105  * - DRAM physical pages pool of 2MB.
2106  * - MMU module.
2107  */
2108 void hl_vm_fini(struct hl_device *hdev)
2109 {
2110         struct hl_vm *vm = &hdev->vm;
2111
2112         if (!vm->init_done)
2113                 return;
2114
2115         /*
2116          * At this point all the contexts should be freed and hence no DRAM
2117          * memory should be in use. Hence the DRAM pool should be freed here.
2118          */
2119         if (kref_put(&vm->dram_pg_pool_refcount, dram_pg_pool_do_release) != 1)
2120                 dev_warn(hdev->dev, "dram_pg_pool was not destroyed on %s\n",
2121                                 __func__);
2122
2123         vm->init_done = false;
2124 }