1// SPDX-License-Identifier: GPL-2.0
2
3/*
4 * Copyright 2016-2022 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#include <linux/vmalloc.h>
15#include <linux/pci-p2pdma.h>
16
17MODULE_IMPORT_NS(DMA_BUF);
18
19#define HL_MMU_DEBUG	0
20
21/* use small pages for supporting non-pow2 (32M/40M/48M) DRAM phys page sizes */
22#define DRAM_POOL_PAGE_SIZE SZ_8M
23
24static int allocate_timestamps_buffers(struct hl_fpriv *hpriv,
25			struct hl_mem_in *args, u64 *handle);
26
27static int set_alloc_page_size(struct hl_device *hdev, struct hl_mem_in *args, u32 *page_size)
28{
29	struct asic_fixed_properties *prop = &hdev->asic_prop;
30	u64 psize;
31
32	/*
33	 * for ASIC that supports setting the allocation page size by user we will address
34	 * user's choice only if it is not 0 (as 0 means taking the default page size)
35	 */
36	if (prop->supports_user_set_page_size && args->alloc.page_size) {
37		psize = args->alloc.page_size;
38
39		if (!is_power_of_2(psize)) {
40			dev_err(hdev->dev, "user page size (%#llx) is not power of 2\n", psize);
41			return -EINVAL;
42		}
43	} else {
44		psize = prop->device_mem_alloc_default_page_size;
45	}
46
47	*page_size = psize;
48
49	return 0;
50}
51
52/*
53 * The va ranges in context object contain a list with the available chunks of
54 * device virtual memory.
55 * There is one range for host allocations and one for DRAM allocations.
56 *
57 * On initialization each range contains one chunk of all of its available
58 * virtual range which is a half of the total device virtual range.
59 *
60 * On each mapping of physical pages, a suitable virtual range chunk (with a
61 * minimum size) is selected from the list. If the chunk size equals the
62 * requested size, the chunk is returned. Otherwise, the chunk is split into
63 * two chunks - one to return as result and a remainder to stay in the list.
64 *
65 * On each Unmapping of a virtual address, the relevant virtual chunk is
66 * returned to the list. The chunk is added to the list and if its edges match
67 * the edges of the adjacent chunks (means a contiguous chunk can be created),
68 * the chunks are merged.
69 *
70 * On finish, the list is checked to have only one chunk of all the relevant
71 * virtual range (which is a half of the device total virtual range).
72 * If not (means not all mappings were unmapped), a warning is printed.
73 */
74
75/*
76 * alloc_device_memory() - allocate device memory.
77 * @ctx: pointer to the context structure.
78 * @args: host parameters containing the requested size.
79 * @ret_handle: result handle.
80 *
81 * This function does the following:
82 * - Allocate the requested size rounded up to 'dram_page_size' pages.
83 * - Return unique handle for later map/unmap/free.
84 */
85static int alloc_device_memory(struct hl_ctx *ctx, struct hl_mem_in *args,
86				u32 *ret_handle)
87{
88	struct hl_device *hdev = ctx->hdev;
89	struct hl_vm *vm = &hdev->vm;
90	struct hl_vm_phys_pg_pack *phys_pg_pack;
91	u64 paddr = 0, total_size, num_pgs, i;
92	u32 num_curr_pgs, page_size;
93	bool contiguous;
94	int handle, rc;
95
96	num_curr_pgs = 0;
97
98	rc = set_alloc_page_size(hdev, args, &page_size);
99	if (rc)
100		return rc;
101
102	num_pgs = DIV_ROUND_UP_ULL(args->alloc.mem_size, page_size);
103	total_size = num_pgs * page_size;
104
105	if (!total_size) {
106		dev_err(hdev->dev, "Cannot allocate 0 bytes\n");
107		return -EINVAL;
108	}
109
110	contiguous = args->flags & HL_MEM_CONTIGUOUS;
111
112	if (contiguous) {
113		if (is_power_of_2(page_size))
114			paddr = (uintptr_t) gen_pool_dma_alloc_align(vm->dram_pg_pool,
115								     total_size, NULL, page_size);
116		else
117			paddr = gen_pool_alloc(vm->dram_pg_pool, total_size);
118		if (!paddr) {
119			dev_err(hdev->dev,
120				"Cannot allocate %llu contiguous pages with total size of %llu\n",
121				num_pgs, total_size);
122			return -ENOMEM;
123		}
124	}
125
126	phys_pg_pack = kzalloc(sizeof(*phys_pg_pack), GFP_KERNEL);
127	if (!phys_pg_pack) {
128		rc = -ENOMEM;
129		goto pages_pack_err;
130	}
131
132	phys_pg_pack->vm_type = VM_TYPE_PHYS_PACK;
133	phys_pg_pack->asid = ctx->asid;
134	phys_pg_pack->npages = num_pgs;
135	phys_pg_pack->page_size = page_size;
136	phys_pg_pack->total_size = total_size;
137	phys_pg_pack->flags = args->flags;
138	phys_pg_pack->contiguous = contiguous;
139
140	phys_pg_pack->pages = kvmalloc_array(num_pgs, sizeof(u64), GFP_KERNEL);
141	if (ZERO_OR_NULL_PTR(phys_pg_pack->pages)) {
142		rc = -ENOMEM;
143		goto pages_arr_err;
144	}
145
146	if (phys_pg_pack->contiguous) {
147		for (i = 0 ; i < num_pgs ; i++)
148			phys_pg_pack->pages[i] = paddr + i * page_size;
149	} else {
150		for (i = 0 ; i < num_pgs ; i++) {
151			if (is_power_of_2(page_size))
152				phys_pg_pack->pages[i] =
153					(uintptr_t)gen_pool_dma_alloc_align(vm->dram_pg_pool,
154									    page_size, NULL,
155									    page_size);
156			else
157				phys_pg_pack->pages[i] = gen_pool_alloc(vm->dram_pg_pool,
158									page_size);
159
160			if (!phys_pg_pack->pages[i]) {
161				dev_err(hdev->dev,
162					"Cannot allocate device memory (out of memory)\n");
163				rc = -ENOMEM;
164				goto page_err;
165			}
166
167			num_curr_pgs++;
168		}
169	}
170
171	spin_lock(&vm->idr_lock);
172	handle = idr_alloc(&vm->phys_pg_pack_handles, phys_pg_pack, 1, 0,
173				GFP_ATOMIC);
174	spin_unlock(&vm->idr_lock);
175
176	if (handle < 0) {
177		dev_err(hdev->dev, "Failed to get handle for page\n");
178		rc = -EFAULT;
179		goto idr_err;
180	}
181
182	for (i = 0 ; i < num_pgs ; i++)
183		kref_get(&vm->dram_pg_pool_refcount);
184
185	phys_pg_pack->handle = handle;
186
187	atomic64_add(phys_pg_pack->total_size, &ctx->dram_phys_mem);
188	atomic64_add(phys_pg_pack->total_size, &hdev->dram_used_mem);
189
190	*ret_handle = handle;
191
192	return 0;
193
194idr_err:
195page_err:
196	if (!phys_pg_pack->contiguous)
197		for (i = 0 ; i < num_curr_pgs ; i++)
198			gen_pool_free(vm->dram_pg_pool, phys_pg_pack->pages[i],
199					page_size);
200
201	kvfree(phys_pg_pack->pages);
202pages_arr_err:
203	kfree(phys_pg_pack);
204pages_pack_err:
205	if (contiguous)
206		gen_pool_free(vm->dram_pg_pool, paddr, total_size);
207
208	return rc;
209}
210
211/**
212 * dma_map_host_va() - DMA mapping of the given host virtual address.
213 * @hdev: habanalabs device structure.
214 * @addr: the host virtual address of the memory area.
215 * @size: the size of the memory area.
216 * @p_userptr: pointer to result userptr structure.
217 *
218 * This function does the following:
219 * - Allocate userptr structure.
220 * - Pin the given host memory using the userptr structure.
221 * - Perform DMA mapping to have the DMA addresses of the pages.
222 */
223static int dma_map_host_va(struct hl_device *hdev, u64 addr, u64 size,
224				struct hl_userptr **p_userptr)
225{
226	struct hl_userptr *userptr;
227	int rc;
228
229	userptr = kzalloc(sizeof(*userptr), GFP_KERNEL);
230	if (!userptr) {
231		rc = -ENOMEM;
232		goto userptr_err;
233	}
234
235	rc = hl_pin_host_memory(hdev, addr, size, userptr);
236	if (rc) {
237		dev_err(hdev->dev, "Failed to pin host memory\n");
238		goto pin_err;
239	}
240
241	userptr->dma_mapped = true;
242	userptr->dir = DMA_BIDIRECTIONAL;
243	userptr->vm_type = VM_TYPE_USERPTR;
244
245	*p_userptr = userptr;
246
247	rc = hdev->asic_funcs->asic_dma_map_sgtable(hdev, userptr->sgt, DMA_BIDIRECTIONAL);
248	if (rc) {
249		dev_err(hdev->dev, "failed to map sgt with DMA region\n");
250		goto dma_map_err;
251	}
252
253	return 0;
254
255dma_map_err:
256	hl_unpin_host_memory(hdev, userptr);
257pin_err:
258	kfree(userptr);
259userptr_err:
260
261	return rc;
262}
263
264/**
265 * dma_unmap_host_va() - DMA unmapping of the given host virtual address.
266 * @hdev: habanalabs device structure.
267 * @userptr: userptr to free.
268 *
269 * This function does the following:
270 * - Unpins the physical pages.
271 * - Frees the userptr structure.
272 */
273static void dma_unmap_host_va(struct hl_device *hdev,
274				struct hl_userptr *userptr)
275{
276	hl_unpin_host_memory(hdev, userptr);
277	kfree(userptr);
278}
279
280/**
281 * dram_pg_pool_do_release() - free DRAM pages pool
282 * @ref: pointer to reference object.
283 *
284 * This function does the following:
285 * - Frees the idr structure of physical pages handles.
286 * - Frees the generic pool of DRAM physical pages.
287 */
288static void dram_pg_pool_do_release(struct kref *ref)
289{
290	struct hl_vm *vm = container_of(ref, struct hl_vm,
291			dram_pg_pool_refcount);
292
293	/*
294	 * free the idr here as only here we know for sure that there are no
295	 * allocated physical pages and hence there are no handles in use
296	 */
297	idr_destroy(&vm->phys_pg_pack_handles);
298	gen_pool_destroy(vm->dram_pg_pool);
299}
300
301/**
302 * free_phys_pg_pack() - free physical page pack.
303 * @hdev: habanalabs device structure.
304 * @phys_pg_pack: physical page pack to free.
305 *
306 * This function does the following:
307 * - For DRAM memory only
308 *   - iterate over the pack, free each physical block structure by
309 *     returning it to the general pool.
310 * - Free the hl_vm_phys_pg_pack structure.
311 */
312static void free_phys_pg_pack(struct hl_device *hdev,
313				struct hl_vm_phys_pg_pack *phys_pg_pack)
314{
315	struct hl_vm *vm = &hdev->vm;
316	u64 i;
317
318	if (phys_pg_pack->created_from_userptr)
319		goto end;
320
321	if (phys_pg_pack->contiguous) {
322		gen_pool_free(vm->dram_pg_pool, phys_pg_pack->pages[0],
323			phys_pg_pack->total_size);
324
325		for (i = 0; i < phys_pg_pack->npages ; i++)
326			kref_put(&vm->dram_pg_pool_refcount,
327				dram_pg_pool_do_release);
328	} else {
329		for (i = 0 ; i < phys_pg_pack->npages ; i++) {
330			gen_pool_free(vm->dram_pg_pool,
331				phys_pg_pack->pages[i],
332				phys_pg_pack->page_size);
333			kref_put(&vm->dram_pg_pool_refcount,
334				dram_pg_pool_do_release);
335		}
336	}
337
338end:
339	kvfree(phys_pg_pack->pages);
340	kfree(phys_pg_pack);
341
342	return;
343}
344
345/**
346 * free_device_memory() - free device memory.
347 * @ctx: pointer to the context structure.
348 * @args: host parameters containing the requested size.
349 *
350 * This function does the following:
351 * - Free the device memory related to the given handle.
352 */
353static int free_device_memory(struct hl_ctx *ctx, struct hl_mem_in *args)
354{
355	struct hl_device *hdev = ctx->hdev;
356	struct hl_vm *vm = &hdev->vm;
357	struct hl_vm_phys_pg_pack *phys_pg_pack;
358	u32 handle = args->free.handle;
359
360	spin_lock(&vm->idr_lock);
361	phys_pg_pack = idr_find(&vm->phys_pg_pack_handles, handle);
362	if (!phys_pg_pack) {
363		spin_unlock(&vm->idr_lock);
364		dev_err(hdev->dev, "free device memory failed, no match for handle %u\n", handle);
365		return -EINVAL;
366	}
367
368	if (atomic_read(&phys_pg_pack->mapping_cnt) > 0) {
369		spin_unlock(&vm->idr_lock);
370		dev_err(hdev->dev, "handle %u is mapped, cannot free\n", handle);
371		return -EINVAL;
372	}
373
374	if (phys_pg_pack->exporting_cnt) {
375		spin_unlock(&vm->idr_lock);
376		dev_dbg(hdev->dev, "handle %u is exported, cannot free\n", handle);
377		return -EINVAL;
378	}
379
380	/* must remove from idr before the freeing of the physical pages as the refcount of the pool
381	 * is also the trigger of the idr destroy
382	 */
383	idr_remove(&vm->phys_pg_pack_handles, handle);
384	spin_unlock(&vm->idr_lock);
385
386	atomic64_sub(phys_pg_pack->total_size, &ctx->dram_phys_mem);
387	atomic64_sub(phys_pg_pack->total_size, &hdev->dram_used_mem);
388
389	free_phys_pg_pack(hdev, phys_pg_pack);
390
391	return 0;
392}
393
394/**
395 * clear_va_list_locked() - free virtual addresses list.
396 * @hdev: habanalabs device structure.
397 * @va_list: list of virtual addresses to free.
398 *
399 * This function does the following:
400 * - Iterate over the list and free each virtual addresses block.
401 *
402 * This function should be called only when va_list lock is taken.
403 */
404static void clear_va_list_locked(struct hl_device *hdev,
405		struct list_head *va_list)
406{
407	struct hl_vm_va_block *va_block, *tmp;
408
409	list_for_each_entry_safe(va_block, tmp, va_list, node) {
410		list_del(&va_block->node);
411		kfree(va_block);
412	}
413}
414
415/**
416 * print_va_list_locked() - print virtual addresses list.
417 * @hdev: habanalabs device structure.
418 * @va_list: list of virtual addresses to print.
419 *
420 * This function does the following:
421 * - Iterate over the list and print each virtual addresses block.
422 *
423 * This function should be called only when va_list lock is taken.
424 */
425static void print_va_list_locked(struct hl_device *hdev,
426		struct list_head *va_list)
427{
428#if HL_MMU_DEBUG
429	struct hl_vm_va_block *va_block;
430
431	dev_dbg(hdev->dev, "print va list:\n");
432
433	list_for_each_entry(va_block, va_list, node)
434		dev_dbg(hdev->dev,
435			"va block, start: 0x%llx, end: 0x%llx, size: %llu\n",
436			va_block->start, va_block->end, va_block->size);
437#endif
438}
439
440/**
441 * merge_va_blocks_locked() - merge a virtual block if possible.
442 * @hdev: pointer to the habanalabs device structure.
443 * @va_list: pointer to the virtual addresses block list.
444 * @va_block: virtual block to merge with adjacent blocks.
445 *
446 * This function does the following:
447 * - Merge the given blocks with the adjacent blocks if their virtual ranges
448 *   create a contiguous virtual range.
449 *
450 * This Function should be called only when va_list lock is taken.
451 */
452static void merge_va_blocks_locked(struct hl_device *hdev,
453		struct list_head *va_list, struct hl_vm_va_block *va_block)
454{
455	struct hl_vm_va_block *prev, *next;
456
457	prev = list_prev_entry(va_block, node);
458	if (&prev->node != va_list && prev->end + 1 == va_block->start) {
459		prev->end = va_block->end;
460		prev->size = prev->end - prev->start;
461		list_del(&va_block->node);
462		kfree(va_block);
463		va_block = prev;
464	}
465
466	next = list_next_entry(va_block, node);
467	if (&next->node != va_list && va_block->end + 1 == next->start) {
468		next->start = va_block->start;
469		next->size = next->end - next->start;
470		list_del(&va_block->node);
471		kfree(va_block);
472	}
473}
474
475/**
476 * add_va_block_locked() - add a virtual block to the virtual addresses list.
477 * @hdev: pointer to the habanalabs device structure.
478 * @va_list: pointer to the virtual addresses block list.
479 * @start: start virtual address.
480 * @end: end virtual address.
481 *
482 * This function does the following:
483 * - Add the given block to the virtual blocks list and merge with other blocks
484 *   if a contiguous virtual block can be created.
485 *
486 * This Function should be called only when va_list lock is taken.
487 */
488static int add_va_block_locked(struct hl_device *hdev,
489		struct list_head *va_list, u64 start, u64 end)
490{
491	struct hl_vm_va_block *va_block, *res = NULL;
492	u64 size = end - start + 1;
493
494	print_va_list_locked(hdev, va_list);
495
496	list_for_each_entry(va_block, va_list, node) {
497		/* TODO: remove upon matureness */
498		if (hl_mem_area_crosses_range(start, size, va_block->start,
499				va_block->end)) {
500			dev_err(hdev->dev,
501				"block crossing ranges at start 0x%llx, end 0x%llx\n",
502				va_block->start, va_block->end);
503			return -EINVAL;
504		}
505
506		if (va_block->end < start)
507			res = va_block;
508	}
509
510	va_block = kmalloc(sizeof(*va_block), GFP_KERNEL);
511	if (!va_block)
512		return -ENOMEM;
513
514	va_block->start = start;
515	va_block->end = end;
516	va_block->size = size;
517
518	if (!res)
519		list_add(&va_block->node, va_list);
520	else
521		list_add(&va_block->node, &res->node);
522
523	merge_va_blocks_locked(hdev, va_list, va_block);
524
525	print_va_list_locked(hdev, va_list);
526
527	return 0;
528}
529
530/**
531 * add_va_block() - wrapper for add_va_block_locked.
532 * @hdev: pointer to the habanalabs device structure.
533 * @va_range: pointer to the virtual addresses range object.
534 * @start: start virtual address.
535 * @end: end virtual address.
536 *
537 * This function does the following:
538 * - Takes the list lock and calls add_va_block_locked.
539 */
540static inline int add_va_block(struct hl_device *hdev,
541		struct hl_va_range *va_range, u64 start, u64 end)
542{
543	int rc;
544
545	mutex_lock(&va_range->lock);
546	rc = add_va_block_locked(hdev, &va_range->list, start, end);
547	mutex_unlock(&va_range->lock);
548
549	return rc;
550}
551
552/**
553 * is_hint_crossing_range() - check if hint address crossing specified reserved.
554 * @range_type: virtual space range type.
555 * @start_addr: start virtual address.
556 * @size: block size.
557 * @prop: asic properties structure to retrieve reserved ranges from.
558 */
559static inline bool is_hint_crossing_range(enum hl_va_range_type range_type,
560		u64 start_addr, u32 size, struct asic_fixed_properties *prop) {
561	bool range_cross;
562
563	if (range_type == HL_VA_RANGE_TYPE_DRAM)
564		range_cross =
565			hl_mem_area_crosses_range(start_addr, size,
566			prop->hints_dram_reserved_va_range.start_addr,
567			prop->hints_dram_reserved_va_range.end_addr);
568	else if (range_type == HL_VA_RANGE_TYPE_HOST)
569		range_cross =
570			hl_mem_area_crosses_range(start_addr,	size,
571			prop->hints_host_reserved_va_range.start_addr,
572			prop->hints_host_reserved_va_range.end_addr);
573	else
574		range_cross =
575			hl_mem_area_crosses_range(start_addr, size,
576			prop->hints_host_hpage_reserved_va_range.start_addr,
577			prop->hints_host_hpage_reserved_va_range.end_addr);
578
579	return range_cross;
580}
581
582/**
583 * get_va_block() - get a virtual block for the given size and alignment.
584 *
585 * @hdev: pointer to the habanalabs device structure.
586 * @va_range: pointer to the virtual addresses range.
587 * @size: requested block size.
588 * @hint_addr: hint for requested address by the user.
589 * @va_block_align: required alignment of the virtual block start address.
590 * @range_type: va range type (host, dram)
591 * @flags: additional memory flags, currently only uses HL_MEM_FORCE_HINT
592 *
593 * This function does the following:
594 * - Iterate on the virtual block list to find a suitable virtual block for the
595 *   given size, hint address and alignment.
596 * - Reserve the requested block and update the list.
597 * - Return the start address of the virtual block.
598 */
599static u64 get_va_block(struct hl_device *hdev,
600				struct hl_va_range *va_range,
601				u64 size, u64 hint_addr, u32 va_block_align,
602				enum hl_va_range_type range_type,
603				u32 flags)
604{
605	struct hl_vm_va_block *va_block, *new_va_block = NULL;
606	struct asic_fixed_properties *prop = &hdev->asic_prop;
607	u64 tmp_hint_addr, valid_start, valid_size, prev_start, prev_end,
608		align_mask, reserved_valid_start = 0, reserved_valid_size = 0,
609		dram_hint_mask = prop->dram_hints_align_mask;
610	bool add_prev = false;
611	bool is_align_pow_2  = is_power_of_2(va_range->page_size);
612	bool is_hint_dram_addr = hl_is_dram_va(hdev, hint_addr);
613	bool force_hint = flags & HL_MEM_FORCE_HINT;
614
615	if (is_align_pow_2)
616		align_mask = ~((u64)va_block_align - 1);
617	else
618		/*
619		 * with non-power-of-2 range we work only with page granularity
620		 * and the start address is page aligned,
621		 * so no need for alignment checking.
622		 */
623		size = DIV_ROUND_UP_ULL(size, va_range->page_size) *
624							va_range->page_size;
625
626	tmp_hint_addr = hint_addr & ~dram_hint_mask;
627
628	/* Check if we need to ignore hint address */
629	if ((is_align_pow_2 && (hint_addr & (va_block_align - 1))) ||
630			(!is_align_pow_2 && is_hint_dram_addr &&
631			do_div(tmp_hint_addr, va_range->page_size))) {
632
633		if (force_hint) {
634			/* Hint must be respected, so here we just fail */
635			dev_err(hdev->dev,
636				"Hint address 0x%llx is not page aligned - cannot be respected\n",
637				hint_addr);
638			return 0;
639		}
640
641		dev_dbg(hdev->dev,
642			"Hint address 0x%llx will be ignored because it is not aligned\n",
643			hint_addr);
644		hint_addr = 0;
645	}
646
647	mutex_lock(&va_range->lock);
648
649	print_va_list_locked(hdev, &va_range->list);
650
651	list_for_each_entry(va_block, &va_range->list, node) {
652		/* Calc the first possible aligned addr */
653		valid_start = va_block->start;
654
655		if (is_align_pow_2 && (valid_start & (va_block_align - 1))) {
656			valid_start &= align_mask;
657			valid_start += va_block_align;
658			if (valid_start > va_block->end)
659				continue;
660		}
661
662		valid_size = va_block->end - valid_start + 1;
663		if (valid_size < size)
664			continue;
665
666		/*
667		 * In case hint address is 0, and hints_range_reservation
668		 * property enabled, then avoid allocating va blocks from the
669		 * range reserved for hint addresses
670		 */
671		if (prop->hints_range_reservation && !hint_addr)
672			if (is_hint_crossing_range(range_type, valid_start,
673					size, prop))
674				continue;
675
676		/* Pick the minimal length block which has the required size */
677		if (!new_va_block || (valid_size < reserved_valid_size)) {
678			new_va_block = va_block;
679			reserved_valid_start = valid_start;
680			reserved_valid_size = valid_size;
681		}
682
683		if (hint_addr && hint_addr >= valid_start &&
684					(hint_addr + size) <= va_block->end) {
685			new_va_block = va_block;
686			reserved_valid_start = hint_addr;
687			reserved_valid_size = valid_size;
688			break;
689		}
690	}
691
692	if (!new_va_block) {
693		dev_err(hdev->dev, "no available va block for size %llu\n",
694								size);
695		goto out;
696	}
697
698	if (force_hint && reserved_valid_start != hint_addr) {
699		/* Hint address must be respected. If we are here - this means
700		 * we could not respect it.
701		 */
702		dev_err(hdev->dev,
703			"Hint address 0x%llx could not be respected\n",
704			hint_addr);
705		reserved_valid_start = 0;
706		goto out;
707	}
708
709	/*
710	 * Check if there is some leftover range due to reserving the new
711	 * va block, then return it to the main virtual addresses list.
712	 */
713	if (reserved_valid_start > new_va_block->start) {
714		prev_start = new_va_block->start;
715		prev_end = reserved_valid_start - 1;
716
717		new_va_block->start = reserved_valid_start;
718		new_va_block->size = reserved_valid_size;
719
720		add_prev = true;
721	}
722
723	if (new_va_block->size > size) {
724		new_va_block->start += size;
725		new_va_block->size = new_va_block->end - new_va_block->start + 1;
726	} else {
727		list_del(&new_va_block->node);
728		kfree(new_va_block);
729	}
730
731	if (add_prev)
732		add_va_block_locked(hdev, &va_range->list, prev_start,
733				prev_end);
734
735	print_va_list_locked(hdev, &va_range->list);
736out:
737	mutex_unlock(&va_range->lock);
738
739	return reserved_valid_start;
740}
741
742/*
743 * hl_reserve_va_block() - reserve a virtual block of a given size.
744 * @hdev: pointer to the habanalabs device structure.
745 * @ctx: current context
746 * @type: virtual addresses range type.
747 * @size: requested block size.
748 * @alignment: required alignment in bytes of the virtual block start address,
749 *             0 means no alignment.
750 *
751 * This function does the following:
752 * - Iterate on the virtual block list to find a suitable virtual block for the
753 *   given size and alignment.
754 * - Reserve the requested block and update the list.
755 * - Return the start address of the virtual block.
756 */
757u64 hl_reserve_va_block(struct hl_device *hdev, struct hl_ctx *ctx,
758		enum hl_va_range_type type, u32 size, u32 alignment)
759{
760	return get_va_block(hdev, ctx->va_range[type], size, 0,
761			max(alignment, ctx->va_range[type]->page_size),
762			type, 0);
763}
764
765/**
766 * hl_get_va_range_type() - get va_range type for the given address and size.
767 * @ctx: context to fetch va_range from.
768 * @address: the start address of the area we want to validate.
769 * @size: the size in bytes of the area we want to validate.
770 * @type: returned va_range type.
771 *
772 * Return: true if the area is inside a valid range, false otherwise.
773 */
774static int hl_get_va_range_type(struct hl_ctx *ctx, u64 address, u64 size,
775			enum hl_va_range_type *type)
776{
777	int i;
778
779	for (i = 0 ; i < HL_VA_RANGE_TYPE_MAX; i++) {
780		if (hl_mem_area_inside_range(address, size,
781				ctx->va_range[i]->start_addr,
782				ctx->va_range[i]->end_addr)) {
783			*type = i;
784			return 0;
785		}
786	}
787
788	return -EINVAL;
789}
790
791/**
792 * hl_unreserve_va_block() - wrapper for add_va_block to unreserve a va block.
793 * @hdev: pointer to the habanalabs device structure
794 * @ctx: pointer to the context structure.
795 * @start_addr: start virtual address.
796 * @size: number of bytes to unreserve.
797 *
798 * This function does the following:
799 * - Takes the list lock and calls add_va_block_locked.
800 */
801int hl_unreserve_va_block(struct hl_device *hdev, struct hl_ctx *ctx,
802		u64 start_addr, u64 size)
803{
804	enum hl_va_range_type type;
805	int rc;
806
807	rc = hl_get_va_range_type(ctx, start_addr, size, &type);
808	if (rc) {
809		dev_err(hdev->dev,
810			"cannot find va_range for va %#llx size %llu",
811			start_addr, size);
812		return rc;
813	}
814
815	rc = add_va_block(hdev, ctx->va_range[type], start_addr,
816						start_addr + size - 1);
817	if (rc)
818		dev_warn(hdev->dev,
819			"add va block failed for vaddr: 0x%llx\n", start_addr);
820
821	return rc;
822}
823
824/**
825 * init_phys_pg_pack_from_userptr() - initialize physical page pack from host
826 *                                    memory
827 * @ctx: pointer to the context structure.
828 * @userptr: userptr to initialize from.
829 * @pphys_pg_pack: result pointer.
830 * @force_regular_page: tell the function to ignore huge page optimization,
831 *                      even if possible. Needed for cases where the device VA
832 *                      is allocated before we know the composition of the
833 *                      physical pages
834 *
835 * This function does the following:
836 * - Pin the physical pages related to the given virtual block.
837 * - Create a physical page pack from the physical pages related to the given
838 *   virtual block.
839 */
840static int init_phys_pg_pack_from_userptr(struct hl_ctx *ctx,
841				struct hl_userptr *userptr,
842				struct hl_vm_phys_pg_pack **pphys_pg_pack,
843				bool force_regular_page)
844{
845	u32 npages, page_size = PAGE_SIZE,
846		huge_page_size = ctx->hdev->asic_prop.pmmu_huge.page_size;
847	u32 pgs_in_huge_page = huge_page_size >> __ffs(page_size);
848	struct hl_vm_phys_pg_pack *phys_pg_pack;
849	bool first = true, is_huge_page_opt;
850	u64 page_mask, total_npages;
851	struct scatterlist *sg;
852	dma_addr_t dma_addr;
853	int rc, i, j;
854
855	phys_pg_pack = kzalloc(sizeof(*phys_pg_pack), GFP_KERNEL);
856	if (!phys_pg_pack)
857		return -ENOMEM;
858
859	phys_pg_pack->vm_type = userptr->vm_type;
860	phys_pg_pack->created_from_userptr = true;
861	phys_pg_pack->asid = ctx->asid;
862	atomic_set(&phys_pg_pack->mapping_cnt, 1);
863
864	is_huge_page_opt = (force_regular_page ? false : true);
865
866	/* Only if all dma_addrs are aligned to 2MB and their
867	 * sizes is at least 2MB, we can use huge page mapping.
868	 * We limit the 2MB optimization to this condition,
869	 * since later on we acquire the related VA range as one
870	 * consecutive block.
871	 */
872	total_npages = 0;
873	for_each_sgtable_dma_sg(userptr->sgt, sg, i) {
874		npages = hl_get_sg_info(sg, &dma_addr);
875
876		total_npages += npages;
877
878		if ((npages % pgs_in_huge_page) ||
879					(dma_addr & (huge_page_size - 1)))
880			is_huge_page_opt = false;
881	}
882
883	if (is_huge_page_opt) {
884		page_size = huge_page_size;
885		do_div(total_npages, pgs_in_huge_page);
886	}
887
888	page_mask = ~(((u64) page_size) - 1);
889
890	phys_pg_pack->pages = kvmalloc_array(total_npages, sizeof(u64),
891						GFP_KERNEL);
892	if (ZERO_OR_NULL_PTR(phys_pg_pack->pages)) {
893		rc = -ENOMEM;
894		goto page_pack_arr_mem_err;
895	}
896
897	phys_pg_pack->npages = total_npages;
898	phys_pg_pack->page_size = page_size;
899	phys_pg_pack->total_size = total_npages * page_size;
900
901	j = 0;
902	for_each_sgtable_dma_sg(userptr->sgt, sg, i) {
903		npages = hl_get_sg_info(sg, &dma_addr);
904
905		/* align down to physical page size and save the offset */
906		if (first) {
907			first = false;
908			phys_pg_pack->offset = dma_addr & (page_size - 1);
909			dma_addr &= page_mask;
910		}
911
912		while (npages) {
913			phys_pg_pack->pages[j++] = dma_addr;
914			dma_addr += page_size;
915
916			if (is_huge_page_opt)
917				npages -= pgs_in_huge_page;
918			else
919				npages--;
920		}
921	}
922
923	*pphys_pg_pack = phys_pg_pack;
924
925	return 0;
926
927page_pack_arr_mem_err:
928	kfree(phys_pg_pack);
929
930	return rc;
931}
932
933/**
934 * map_phys_pg_pack() - maps the physical page pack..
935 * @ctx: pointer to the context structure.
936 * @vaddr: start address of the virtual area to map from.
937 * @phys_pg_pack: the pack of physical pages to map to.
938 *
939 * This function does the following:
940 * - Maps each chunk of virtual memory to matching physical chunk.
941 * - Stores number of successful mappings in the given argument.
942 * - Returns 0 on success, error code otherwise.
943 */
944static int map_phys_pg_pack(struct hl_ctx *ctx, u64 vaddr,
945				struct hl_vm_phys_pg_pack *phys_pg_pack)
946{
947	struct hl_device *hdev = ctx->hdev;
948	u64 next_vaddr = vaddr, paddr, mapped_pg_cnt = 0, i;
949	u32 page_size = phys_pg_pack->page_size;
950	int rc = 0;
951	bool is_host_addr;
952
953	for (i = 0 ; i < phys_pg_pack->npages ; i++) {
954		paddr = phys_pg_pack->pages[i];
955
956		rc = hl_mmu_map_page(ctx, next_vaddr, paddr, page_size,
957				(i + 1) == phys_pg_pack->npages);
958		if (rc) {
959			dev_err(hdev->dev,
960				"map failed for handle %u, npages: %llu, mapped: %llu",
961				phys_pg_pack->handle, phys_pg_pack->npages,
962				mapped_pg_cnt);
963			goto err;
964		}
965
966		mapped_pg_cnt++;
967		next_vaddr += page_size;
968	}
969
970	return 0;
971
972err:
973	is_host_addr = !hl_is_dram_va(hdev, vaddr);
974
975	next_vaddr = vaddr;
976	for (i = 0 ; i < mapped_pg_cnt ; i++) {
977		if (hl_mmu_unmap_page(ctx, next_vaddr, page_size,
978					(i + 1) == mapped_pg_cnt))
979			dev_warn_ratelimited(hdev->dev,
980				"failed to unmap handle %u, va: 0x%llx, pa: 0x%llx, page size: %u\n",
981					phys_pg_pack->handle, next_vaddr,
982					phys_pg_pack->pages[i], page_size);
983
984		next_vaddr += page_size;
985
986		/*
987		 * unmapping on Palladium can be really long, so avoid a CPU
988		 * soft lockup bug by sleeping a little between unmapping pages
989		 *
990		 * In addition, on host num of pages could be huge,
991		 * because page size could be 4KB, so when unmapping host
992		 * pages sleep every 32K pages to avoid soft lockup
993		 */
994		if (hdev->pldm || (is_host_addr && (i & 0x7FFF) == 0))
995			usleep_range(50, 200);
996	}
997
998	return rc;
999}
1000
1001/**
1002 * unmap_phys_pg_pack() - unmaps the physical page pack.
1003 * @ctx: pointer to the context structure.
1004 * @vaddr: start address of the virtual area to unmap.
1005 * @phys_pg_pack: the pack of physical pages to unmap.
1006 */
1007static void unmap_phys_pg_pack(struct hl_ctx *ctx, u64 vaddr,
1008				struct hl_vm_phys_pg_pack *phys_pg_pack)
1009{
1010	struct hl_device *hdev = ctx->hdev;
1011	u64 next_vaddr, i;
1012	bool is_host_addr;
1013	u32 page_size;
1014
1015	is_host_addr = !hl_is_dram_va(hdev, vaddr);
1016	page_size = phys_pg_pack->page_size;
1017	next_vaddr = vaddr;
1018
1019	for (i = 0 ; i < phys_pg_pack->npages ; i++, next_vaddr += page_size) {
1020		if (hl_mmu_unmap_page(ctx, next_vaddr, page_size,
1021				       (i + 1) == phys_pg_pack->npages))
1022			dev_warn_ratelimited(hdev->dev,
1023			"unmap failed for vaddr: 0x%llx\n", next_vaddr);
1024
1025		/*
1026		 * unmapping on Palladium can be really long, so avoid a CPU
1027		 * soft lockup bug by sleeping a little between unmapping pages
1028		 *
1029		 * In addition, on host num of pages could be huge,
1030		 * because page size could be 4KB, so when unmapping host
1031		 * pages sleep every 32K pages to avoid soft lockup
1032		 */
1033		if (hdev->pldm || (is_host_addr && (i & 0x7FFF) == 0))
1034			usleep_range(50, 200);
1035	}
1036}
1037
1038static int get_paddr_from_handle(struct hl_ctx *ctx, struct hl_mem_in *args,
1039					u64 *paddr)
1040{
1041	struct hl_device *hdev = ctx->hdev;
1042	struct hl_vm *vm = &hdev->vm;
1043	struct hl_vm_phys_pg_pack *phys_pg_pack;
1044	u32 handle;
1045
1046	handle = lower_32_bits(args->map_device.handle);
1047	spin_lock(&vm->idr_lock);
1048	phys_pg_pack = idr_find(&vm->phys_pg_pack_handles, handle);
1049	if (!phys_pg_pack) {
1050		spin_unlock(&vm->idr_lock);
1051		dev_err(hdev->dev, "no match for handle %u\n", handle);
1052		return -EINVAL;
1053	}
1054
1055	*paddr = phys_pg_pack->pages[0];
1056
1057	spin_unlock(&vm->idr_lock);
1058
1059	return 0;
1060}
1061
1062/**
1063 * map_device_va() - map the given memory.
1064 * @ctx: pointer to the context structure.
1065 * @args: host parameters with handle/host virtual address.
1066 * @device_addr: pointer to result device virtual address.
1067 *
1068 * This function does the following:
1069 * - If given a physical device memory handle, map to a device virtual block
1070 *   and return the start address of this block.
1071 * - If given a host virtual address and size, find the related physical pages,
1072 *   map a device virtual block to this pages and return the start address of
1073 *   this block.
1074 */
1075static int map_device_va(struct hl_ctx *ctx, struct hl_mem_in *args, u64 *device_addr)
1076{
1077	struct hl_vm_phys_pg_pack *phys_pg_pack;
1078	enum hl_va_range_type va_range_type = 0;
1079	struct hl_device *hdev = ctx->hdev;
1080	struct hl_userptr *userptr = NULL;
1081	u32 handle = 0, va_block_align;
1082	struct hl_vm_hash_node *hnode;
1083	struct hl_vm *vm = &hdev->vm;
1084	struct hl_va_range *va_range;
1085	bool is_userptr, do_prefetch;
1086	u64 ret_vaddr, hint_addr;
1087	enum vm_type *vm_type;
1088	int rc;
1089
1090	/* set map flags */
1091	is_userptr = args->flags & HL_MEM_USERPTR;
1092	do_prefetch = hdev->supports_mmu_prefetch && (args->flags & HL_MEM_PREFETCH);
1093
1094	/* Assume failure */
1095	*device_addr = 0;
1096
1097	if (is_userptr) {
1098		u64 addr = args->map_host.host_virt_addr,
1099			size = args->map_host.mem_size;
1100		u32 page_size = hdev->asic_prop.pmmu.page_size,
1101			huge_page_size = hdev->asic_prop.pmmu_huge.page_size;
1102
1103		rc = dma_map_host_va(hdev, addr, size, &userptr);
1104		if (rc) {
1105			dev_err(hdev->dev, "failed to get userptr from va\n");
1106			return rc;
1107		}
1108
1109		rc = init_phys_pg_pack_from_userptr(ctx, userptr,
1110				&phys_pg_pack, false);
1111		if (rc) {
1112			dev_err(hdev->dev,
1113				"unable to init page pack for vaddr 0x%llx\n",
1114				addr);
1115			goto init_page_pack_err;
1116		}
1117
1118		vm_type = (enum vm_type *) userptr;
1119		hint_addr = args->map_host.hint_addr;
1120		handle = phys_pg_pack->handle;
1121
1122		/* get required alignment */
1123		if (phys_pg_pack->page_size == page_size) {
1124			va_range = ctx->va_range[HL_VA_RANGE_TYPE_HOST];
1125			va_range_type = HL_VA_RANGE_TYPE_HOST;
1126			/*
1127			 * huge page alignment may be needed in case of regular
1128			 * page mapping, depending on the host VA alignment
1129			 */
1130			if (addr & (huge_page_size - 1))
1131				va_block_align = page_size;
1132			else
1133				va_block_align = huge_page_size;
1134		} else {
1135			/*
1136			 * huge page alignment is needed in case of huge page
1137			 * mapping
1138			 */
1139			va_range = ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE];
1140			va_range_type = HL_VA_RANGE_TYPE_HOST_HUGE;
1141			va_block_align = huge_page_size;
1142		}
1143	} else {
1144		handle = lower_32_bits(args->map_device.handle);
1145
1146		spin_lock(&vm->idr_lock);
1147		phys_pg_pack = idr_find(&vm->phys_pg_pack_handles, handle);
1148		if (!phys_pg_pack) {
1149			spin_unlock(&vm->idr_lock);
1150			dev_err(hdev->dev,
1151				"no match for handle %u\n", handle);
1152			return -EINVAL;
1153		}
1154
1155		/* increment now to avoid freeing device memory while mapping */
1156		atomic_inc(&phys_pg_pack->mapping_cnt);
1157
1158		spin_unlock(&vm->idr_lock);
1159
1160		vm_type = (enum vm_type *) phys_pg_pack;
1161
1162		hint_addr = args->map_device.hint_addr;
1163
1164		/* DRAM VA alignment is the same as the MMU page size */
1165		va_range = ctx->va_range[HL_VA_RANGE_TYPE_DRAM];
1166		va_range_type = HL_VA_RANGE_TYPE_DRAM;
1167		va_block_align = hdev->asic_prop.dmmu.page_size;
1168	}
1169
1170	/*
1171	 * relevant for mapping device physical memory only, as host memory is
1172	 * implicitly shared
1173	 */
1174	if (!is_userptr && !(phys_pg_pack->flags & HL_MEM_SHARED) &&
1175			phys_pg_pack->asid != ctx->asid) {
1176		dev_err(hdev->dev,
1177			"Failed to map memory, handle %u is not shared\n",
1178			handle);
1179		rc = -EPERM;
1180		goto shared_err;
1181	}
1182
1183	hnode = kzalloc(sizeof(*hnode), GFP_KERNEL);
1184	if (!hnode) {
1185		rc = -ENOMEM;
1186		goto hnode_err;
1187	}
1188
1189	if (hint_addr && phys_pg_pack->offset) {
1190		if (args->flags & HL_MEM_FORCE_HINT) {
1191			/* Fail if hint must be respected but it can't be */
1192			dev_err(hdev->dev,
1193				"Hint address 0x%llx cannot be respected because source memory is not aligned 0x%x\n",
1194				hint_addr, phys_pg_pack->offset);
1195			rc = -EINVAL;
1196			goto va_block_err;
1197		}
1198		dev_dbg(hdev->dev,
1199			"Hint address 0x%llx will be ignored because source memory is not aligned 0x%x\n",
1200			hint_addr, phys_pg_pack->offset);
1201	}
1202
1203	ret_vaddr = get_va_block(hdev, va_range, phys_pg_pack->total_size,
1204					hint_addr, va_block_align,
1205					va_range_type, args->flags);
1206	if (!ret_vaddr) {
1207		dev_err(hdev->dev, "no available va block for handle %u\n",
1208				handle);
1209		rc = -ENOMEM;
1210		goto va_block_err;
1211	}
1212
1213	mutex_lock(&ctx->mmu_lock);
1214
1215	rc = map_phys_pg_pack(ctx, ret_vaddr, phys_pg_pack);
1216	if (rc) {
1217		dev_err(hdev->dev, "mapping page pack failed for handle %u\n", handle);
1218		mutex_unlock(&ctx->mmu_lock);
1219		goto map_err;
1220	}
1221
1222	rc = hl_mmu_invalidate_cache_range(hdev, false, *vm_type | MMU_OP_SKIP_LOW_CACHE_INV,
1223				ctx->asid, ret_vaddr, phys_pg_pack->total_size);
1224	mutex_unlock(&ctx->mmu_lock);
1225	if (rc)
1226		goto map_err;
1227
1228	/*
1229	 * prefetch is done upon user's request. it is performed in WQ as and so can
1230	 * be outside the MMU lock. the operation itself is already protected by the mmu lock
1231	 */
1232	if (do_prefetch) {
1233		rc = hl_mmu_prefetch_cache_range(ctx, *vm_type, ctx->asid, ret_vaddr,
1234							phys_pg_pack->total_size);
1235		if (rc)
1236			goto map_err;
1237	}
1238
1239	ret_vaddr += phys_pg_pack->offset;
1240
1241	hnode->ptr = vm_type;
1242	hnode->vaddr = ret_vaddr;
1243
1244	mutex_lock(&ctx->mem_hash_lock);
1245	hash_add(ctx->mem_hash, &hnode->node, ret_vaddr);
1246	mutex_unlock(&ctx->mem_hash_lock);
1247
1248	*device_addr = ret_vaddr;
1249
1250	if (is_userptr)
1251		free_phys_pg_pack(hdev, phys_pg_pack);
1252
1253	return rc;
1254
1255map_err:
1256	if (add_va_block(hdev, va_range, ret_vaddr,
1257				ret_vaddr + phys_pg_pack->total_size - 1))
1258		dev_warn(hdev->dev,
1259			"release va block failed for handle 0x%x, vaddr: 0x%llx\n",
1260				handle, ret_vaddr);
1261
1262va_block_err:
1263	kfree(hnode);
1264hnode_err:
1265shared_err:
1266	atomic_dec(&phys_pg_pack->mapping_cnt);
1267	if (is_userptr)
1268		free_phys_pg_pack(hdev, phys_pg_pack);
1269init_page_pack_err:
1270	if (is_userptr)
1271		dma_unmap_host_va(hdev, userptr);
1272
1273	return rc;
1274}
1275
1276/**
1277 * unmap_device_va() - unmap the given device virtual address.
1278 * @ctx: pointer to the context structure.
1279 * @args: host parameters with device virtual address to unmap.
1280 * @ctx_free: true if in context free flow, false otherwise.
1281 *
1282 * This function does the following:
1283 * - unmap the physical pages related to the given virtual address.
1284 * - return the device virtual block to the virtual block list.
1285 */
1286static int unmap_device_va(struct hl_ctx *ctx, struct hl_mem_in *args,
1287				bool ctx_free)
1288{
1289	struct hl_vm_phys_pg_pack *phys_pg_pack = NULL;
1290	u64 vaddr = args->unmap.device_virt_addr;
1291	struct hl_vm_hash_node *hnode = NULL;
1292	struct asic_fixed_properties *prop;
1293	struct hl_device *hdev = ctx->hdev;
1294	struct hl_userptr *userptr = NULL;
1295	struct hl_va_range *va_range;
1296	enum vm_type *vm_type;
1297	bool is_userptr;
1298	int rc = 0;
1299
1300	prop = &hdev->asic_prop;
1301
1302	/* protect from double entrance */
1303	mutex_lock(&ctx->mem_hash_lock);
1304	hash_for_each_possible(ctx->mem_hash, hnode, node, (unsigned long)vaddr)
1305		if (vaddr == hnode->vaddr)
1306			break;
1307
1308	if (!hnode) {
1309		mutex_unlock(&ctx->mem_hash_lock);
1310		dev_err(hdev->dev,
1311			"unmap failed, no mem hnode for vaddr 0x%llx\n",
1312			vaddr);
1313		return -EINVAL;
1314	}
1315
1316	hash_del(&hnode->node);
1317	mutex_unlock(&ctx->mem_hash_lock);
1318
1319	vm_type = hnode->ptr;
1320
1321	if (*vm_type == VM_TYPE_USERPTR) {
1322		is_userptr = true;
1323		userptr = hnode->ptr;
1324
1325		rc = init_phys_pg_pack_from_userptr(ctx, userptr, &phys_pg_pack,
1326							false);
1327		if (rc) {
1328			dev_err(hdev->dev,
1329				"unable to init page pack for vaddr 0x%llx\n",
1330				vaddr);
1331			goto vm_type_err;
1332		}
1333
1334		if (phys_pg_pack->page_size ==
1335					hdev->asic_prop.pmmu.page_size)
1336			va_range = ctx->va_range[HL_VA_RANGE_TYPE_HOST];
1337		else
1338			va_range = ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE];
1339	} else if (*vm_type == VM_TYPE_PHYS_PACK) {
1340		is_userptr = false;
1341		va_range = ctx->va_range[HL_VA_RANGE_TYPE_DRAM];
1342		phys_pg_pack = hnode->ptr;
1343	} else {
1344		dev_warn(hdev->dev,
1345			"unmap failed, unknown vm desc for vaddr 0x%llx\n",
1346				vaddr);
1347		rc = -EFAULT;
1348		goto vm_type_err;
1349	}
1350
1351	if (atomic_read(&phys_pg_pack->mapping_cnt) == 0) {
1352		dev_err(hdev->dev, "vaddr 0x%llx is not mapped\n", vaddr);
1353		rc = -EINVAL;
1354		goto mapping_cnt_err;
1355	}
1356
1357	if (!is_userptr && !is_power_of_2(phys_pg_pack->page_size))
1358		vaddr = prop->dram_base_address +
1359			DIV_ROUND_DOWN_ULL(vaddr - prop->dram_base_address,
1360						phys_pg_pack->page_size) *
1361							phys_pg_pack->page_size;
1362	else
1363		vaddr &= ~(((u64) phys_pg_pack->page_size) - 1);
1364
1365	mutex_lock(&ctx->mmu_lock);
1366
1367	unmap_phys_pg_pack(ctx, vaddr, phys_pg_pack);
1368
1369	/*
1370	 * During context free this function is called in a loop to clean all
1371	 * the context mappings. Hence the cache invalidation can be called once
1372	 * at the loop end rather than for each iteration
1373	 */
1374	if (!ctx_free)
1375		rc = hl_mmu_invalidate_cache_range(hdev, true, *vm_type, ctx->asid, vaddr,
1376							phys_pg_pack->total_size);
1377
1378	mutex_unlock(&ctx->mmu_lock);
1379
1380	/*
1381	 * If the context is closing we don't need to check for the MMU cache
1382	 * invalidation return code and update the VA free list as in this flow
1383	 * we invalidate the MMU cache outside of this unmap function and the VA
1384	 * free list will be freed anyway.
1385	 */
1386	if (!ctx_free) {
1387		int tmp_rc;
1388
1389		tmp_rc = add_va_block(hdev, va_range, vaddr,
1390					vaddr + phys_pg_pack->total_size - 1);
1391		if (tmp_rc) {
1392			dev_warn(hdev->dev,
1393					"add va block failed for vaddr: 0x%llx\n",
1394					vaddr);
1395			if (!rc)
1396				rc = tmp_rc;
1397		}
1398	}
1399
1400	atomic_dec(&phys_pg_pack->mapping_cnt);
1401	kfree(hnode);
1402
1403	if (is_userptr) {
1404		free_phys_pg_pack(hdev, phys_pg_pack);
1405		dma_unmap_host_va(hdev, userptr);
1406	}
1407
1408	return rc;
1409
1410mapping_cnt_err:
1411	if (is_userptr)
1412		free_phys_pg_pack(hdev, phys_pg_pack);
1413vm_type_err:
1414	mutex_lock(&ctx->mem_hash_lock);
1415	hash_add(ctx->mem_hash, &hnode->node, vaddr);
1416	mutex_unlock(&ctx->mem_hash_lock);
1417
1418	return rc;
1419}
1420
1421static int map_block(struct hl_device *hdev, u64 address, u64 *handle,
1422			u32 *size)
1423{
1424	u32 block_id = 0;
1425	int rc;
1426
1427	rc = hdev->asic_funcs->get_hw_block_id(hdev, address, size, &block_id);
1428
1429	*handle = block_id | HL_MMAP_TYPE_BLOCK;
1430	*handle <<= PAGE_SHIFT;
1431
1432	return rc;
1433}
1434
1435static void hw_block_vm_close(struct vm_area_struct *vma)
1436{
1437	struct hl_vm_hw_block_list_node *lnode =
1438		(struct hl_vm_hw_block_list_node *) vma->vm_private_data;
1439	struct hl_ctx *ctx = lnode->ctx;
1440
1441	mutex_lock(&ctx->hw_block_list_lock);
1442	list_del(&lnode->node);
1443	mutex_unlock(&ctx->hw_block_list_lock);
1444	hl_ctx_put(ctx);
1445	kfree(lnode);
1446	vma->vm_private_data = NULL;
1447}
1448
1449static const struct vm_operations_struct hw_block_vm_ops = {
1450	.close = hw_block_vm_close
1451};
1452
1453/**
1454 * hl_hw_block_mmap() - mmap a hw block to user.
1455 * @hpriv: pointer to the private data of the fd
1456 * @vma: pointer to vm_area_struct of the process
1457 *
1458 * Driver increments context reference for every HW block mapped in order
1459 * to prevent user from closing FD without unmapping first
1460 */
1461int hl_hw_block_mmap(struct hl_fpriv *hpriv, struct vm_area_struct *vma)
1462{
1463	struct hl_vm_hw_block_list_node *lnode;
1464	struct hl_device *hdev = hpriv->hdev;
1465	struct hl_ctx *ctx = hpriv->ctx;
1466	u32 block_id, block_size;
1467	int rc;
1468
1469	/* We use the page offset to hold the block id and thus we need to clear
1470	 * it before doing the mmap itself
1471	 */
1472	block_id = vma->vm_pgoff;
1473	vma->vm_pgoff = 0;
1474
1475	/* Driver only allows mapping of a complete HW block */
1476	block_size = vma->vm_end - vma->vm_start;
1477
1478	if (!access_ok((void __user *) (uintptr_t) vma->vm_start, block_size)) {
1479		dev_err(hdev->dev,
1480			"user pointer is invalid - 0x%lx\n",
1481			vma->vm_start);
1482
1483		return -EINVAL;
1484	}
1485
1486	lnode = kzalloc(sizeof(*lnode), GFP_KERNEL);
1487	if (!lnode)
1488		return -ENOMEM;
1489
1490	vma->vm_ops = &hw_block_vm_ops;
1491	vma->vm_private_data = lnode;
1492
1493	hl_ctx_get(ctx);
1494
1495	rc = hdev->asic_funcs->hw_block_mmap(hdev, vma, block_id, block_size);
1496	if (rc) {
1497		hl_ctx_put(ctx);
1498		kfree(lnode);
1499		return rc;
1500	}
1501
1502	lnode->ctx = ctx;
1503	lnode->vaddr = vma->vm_start;
1504	lnode->size = block_size;
1505	lnode->id = block_id;
1506
1507	mutex_lock(&ctx->hw_block_list_lock);
1508	list_add_tail(&lnode->node, &ctx->hw_block_mem_list);
1509	mutex_unlock(&ctx->hw_block_list_lock);
1510
1511	vma->vm_pgoff = block_id;
1512
1513	return 0;
1514}
1515
1516static int set_dma_sg(struct scatterlist *sg, u64 bar_address, u64 chunk_size,
1517			struct device *dev, enum dma_data_direction dir)
1518{
1519	dma_addr_t addr;
1520	int rc;
1521
1522	addr = dma_map_resource(dev, bar_address, chunk_size, dir,
1523				DMA_ATTR_SKIP_CPU_SYNC);
1524	rc = dma_mapping_error(dev, addr);
1525	if (rc)
1526		return rc;
1527
1528	sg_set_page(sg, NULL, chunk_size, 0);
1529	sg_dma_address(sg) = addr;
1530	sg_dma_len(sg) = chunk_size;
1531
1532	return 0;
1533}
1534
1535static struct sg_table *alloc_sgt_from_device_pages(struct hl_device *hdev, u64 *pages, u64 npages,
1536						u64 page_size, struct device *dev,
1537						enum dma_data_direction dir)
1538{
1539	u64 chunk_size, bar_address, dma_max_seg_size;
1540	struct asic_fixed_properties *prop;
1541	int rc, i, j, nents, cur_page;
1542	struct scatterlist *sg;
1543	struct sg_table *sgt;
1544
1545	prop = &hdev->asic_prop;
1546
1547	dma_max_seg_size = dma_get_max_seg_size(dev);
1548
1549	/* We would like to align the max segment size to PAGE_SIZE, so the
1550	 * SGL will contain aligned addresses that can be easily mapped to
1551	 * an MMU
1552	 */
1553	dma_max_seg_size = ALIGN_DOWN(dma_max_seg_size, PAGE_SIZE);
1554	if (dma_max_seg_size < PAGE_SIZE) {
1555		dev_err_ratelimited(hdev->dev,
1556				"dma_max_seg_size %llu can't be smaller than PAGE_SIZE\n",
1557				dma_max_seg_size);
1558		return ERR_PTR(-EINVAL);
1559	}
1560
1561	sgt = kzalloc(sizeof(*sgt), GFP_KERNEL);
1562	if (!sgt)
1563		return ERR_PTR(-ENOMEM);
1564
1565	/* If the size of each page is larger than the dma max segment size,
1566	 * then we can't combine pages and the number of entries in the SGL
1567	 * will just be the
1568	 * <number of pages> * <chunks of max segment size in each page>
1569	 */
1570	if (page_size > dma_max_seg_size)
1571		nents = npages * DIV_ROUND_UP_ULL(page_size, dma_max_seg_size);
1572	else
1573		/* Get number of non-contiguous chunks */
1574		for (i = 1, nents = 1, chunk_size = page_size ; i < npages ; i++) {
1575			if (pages[i - 1] + page_size != pages[i] ||
1576					chunk_size + page_size > dma_max_seg_size) {
1577				nents++;
1578				chunk_size = page_size;
1579				continue;
1580			}
1581
1582			chunk_size += page_size;
1583		}
1584
1585	rc = sg_alloc_table(sgt, nents, GFP_KERNEL | __GFP_ZERO);
1586	if (rc)
1587		goto error_free;
1588
1589	cur_page = 0;
1590
1591	if (page_size > dma_max_seg_size) {
1592		u64 size_left, cur_device_address = 0;
1593
1594		size_left = page_size;
1595
1596		/* Need to split each page into the number of chunks of
1597		 * dma_max_seg_size
1598		 */
1599		for_each_sgtable_dma_sg(sgt, sg, i) {
1600			if (size_left == page_size)
1601				cur_device_address =
1602					pages[cur_page] - prop->dram_base_address;
1603			else
1604				cur_device_address += dma_max_seg_size;
1605
1606			chunk_size = min(size_left, dma_max_seg_size);
1607
1608			bar_address = hdev->dram_pci_bar_start + cur_device_address;
1609
1610			rc = set_dma_sg(sg, bar_address, chunk_size, dev, dir);
1611			if (rc)
1612				goto error_unmap;
1613
1614			if (size_left > dma_max_seg_size) {
1615				size_left -= dma_max_seg_size;
1616			} else {
1617				cur_page++;
1618				size_left = page_size;
1619			}
1620		}
1621	} else {
1622		/* Merge pages and put them into the scatterlist */
1623		for_each_sgtable_dma_sg(sgt, sg, i) {
1624			chunk_size = page_size;
1625			for (j = cur_page + 1 ; j < npages ; j++) {
1626				if (pages[j - 1] + page_size != pages[j] ||
1627						chunk_size + page_size > dma_max_seg_size)
1628					break;
1629
1630				chunk_size += page_size;
1631			}
1632
1633			bar_address = hdev->dram_pci_bar_start +
1634					(pages[cur_page] - prop->dram_base_address);
1635
1636			rc = set_dma_sg(sg, bar_address, chunk_size, dev, dir);
1637			if (rc)
1638				goto error_unmap;
1639
1640			cur_page = j;
1641		}
1642	}
1643
1644	/* Because we are not going to include a CPU list we want to have some
1645	 * chance that other users will detect this by setting the orig_nents
1646	 * to 0 and using only nents (length of DMA list) when going over the
1647	 * sgl
1648	 */
1649	sgt->orig_nents = 0;
1650
1651	return sgt;
1652
1653error_unmap:
1654	for_each_sgtable_dma_sg(sgt, sg, i) {
1655		if (!sg_dma_len(sg))
1656			continue;
1657
1658		dma_unmap_resource(dev, sg_dma_address(sg),
1659					sg_dma_len(sg), dir,
1660					DMA_ATTR_SKIP_CPU_SYNC);
1661	}
1662
1663	sg_free_table(sgt);
1664
1665error_free:
1666	kfree(sgt);
1667	return ERR_PTR(rc);
1668}
1669
1670static int hl_dmabuf_attach(struct dma_buf *dmabuf,
1671				struct dma_buf_attachment *attachment)
1672{
1673	struct hl_dmabuf_priv *hl_dmabuf;
1674	struct hl_device *hdev;
1675	int rc;
1676
1677	hl_dmabuf = dmabuf->priv;
1678	hdev = hl_dmabuf->ctx->hdev;
1679
1680	rc = pci_p2pdma_distance_many(hdev->pdev, &attachment->dev, 1, true);
1681
1682	if (rc < 0)
1683		attachment->peer2peer = false;
1684	return 0;
1685}
1686
1687static struct sg_table *hl_map_dmabuf(struct dma_buf_attachment *attachment,
1688					enum dma_data_direction dir)
1689{
1690	struct dma_buf *dma_buf = attachment->dmabuf;
1691	struct hl_vm_phys_pg_pack *phys_pg_pack;
1692	struct hl_dmabuf_priv *hl_dmabuf;
1693	struct hl_device *hdev;
1694	struct sg_table *sgt;
1695
1696	hl_dmabuf = dma_buf->priv;
1697	hdev = hl_dmabuf->ctx->hdev;
1698	phys_pg_pack = hl_dmabuf->phys_pg_pack;
1699
1700	if (!attachment->peer2peer) {
1701		dev_dbg(hdev->dev, "Failed to map dmabuf because p2p is disabled\n");
1702		return ERR_PTR(-EPERM);
1703	}
1704
1705	if (phys_pg_pack)
1706		sgt = alloc_sgt_from_device_pages(hdev,
1707						phys_pg_pack->pages,
1708						phys_pg_pack->npages,
1709						phys_pg_pack->page_size,
1710						attachment->dev,
1711						dir);
1712	else
1713		sgt = alloc_sgt_from_device_pages(hdev,
1714						&hl_dmabuf->device_address,
1715						1,
1716						hl_dmabuf->dmabuf->size,
1717						attachment->dev,
1718						dir);
1719
1720	if (IS_ERR(sgt))
1721		dev_err(hdev->dev, "failed (%ld) to initialize sgt for dmabuf\n", PTR_ERR(sgt));
1722
1723	return sgt;
1724}
1725
1726static void hl_unmap_dmabuf(struct dma_buf_attachment *attachment,
1727				  struct sg_table *sgt,
1728				  enum dma_data_direction dir)
1729{
1730	struct scatterlist *sg;
1731	int i;
1732
1733	/* The memory behind the dma-buf has *always* resided on the device itself, i.e. it lives
1734	 * only in the 'device' domain (after all, it maps a PCI bar address which points to the
1735	 * device memory).
1736	 *
1737	 * Therefore, it was never in the 'CPU' domain and hence, there is no need to perform
1738	 * a sync of the memory to the CPU's cache, as it never resided inside that cache.
1739	 */
1740	for_each_sgtable_dma_sg(sgt, sg, i)
1741		dma_unmap_resource(attachment->dev, sg_dma_address(sg),
1742					sg_dma_len(sg), dir,
1743					DMA_ATTR_SKIP_CPU_SYNC);
1744
1745	/* Need to restore orig_nents because sg_free_table use that field */
1746	sgt->orig_nents = sgt->nents;
1747	sg_free_table(sgt);
1748	kfree(sgt);
1749}
1750
1751static void hl_release_dmabuf(struct dma_buf *dmabuf)
1752{
1753	struct hl_dmabuf_priv *hl_dmabuf = dmabuf->priv;
1754	struct hl_ctx *ctx = hl_dmabuf->ctx;
1755	struct hl_device *hdev = ctx->hdev;
1756	struct hl_vm *vm = &hdev->vm;
1757
1758	if (hl_dmabuf->phys_pg_pack) {
1759		spin_lock(&vm->idr_lock);
1760		hl_dmabuf->phys_pg_pack->exporting_cnt--;
1761		spin_unlock(&vm->idr_lock);
1762	}
1763
1764	hl_ctx_put(hl_dmabuf->ctx);
1765
1766	kfree(hl_dmabuf);
1767}
1768
1769static const struct dma_buf_ops habanalabs_dmabuf_ops = {
1770	.attach = hl_dmabuf_attach,
1771	.map_dma_buf = hl_map_dmabuf,
1772	.unmap_dma_buf = hl_unmap_dmabuf,
1773	.release = hl_release_dmabuf,
1774};
1775
1776static int export_dmabuf_common(struct hl_ctx *ctx,
1777				struct hl_dmabuf_priv *hl_dmabuf,
1778				u64 total_size, int flags, int *dmabuf_fd)
1779{
1780	DEFINE_DMA_BUF_EXPORT_INFO(exp_info);
1781	struct hl_device *hdev = ctx->hdev;
1782	int rc, fd;
1783
1784	exp_info.ops = &habanalabs_dmabuf_ops;
1785	exp_info.size = total_size;
1786	exp_info.flags = flags;
1787	exp_info.priv = hl_dmabuf;
1788
1789	hl_dmabuf->dmabuf = dma_buf_export(&exp_info);
1790	if (IS_ERR(hl_dmabuf->dmabuf)) {
1791		dev_err(hdev->dev, "failed to export dma-buf\n");
1792		return PTR_ERR(hl_dmabuf->dmabuf);
1793	}
1794
1795	fd = dma_buf_fd(hl_dmabuf->dmabuf, flags);
1796	if (fd < 0) {
1797		dev_err(hdev->dev, "failed to get a file descriptor for a dma-buf\n");
1798		rc = fd;
1799		goto err_dma_buf_put;
1800	}
1801
1802	hl_dmabuf->ctx = ctx;
1803	hl_ctx_get(hl_dmabuf->ctx);
1804
1805	*dmabuf_fd = fd;
1806
1807	return 0;
1808
1809err_dma_buf_put:
1810	dma_buf_put(hl_dmabuf->dmabuf);
1811	return rc;
1812}
1813
1814/**
1815 * export_dmabuf_from_addr() - export a dma-buf object for the given memory
1816 *                             address and size.
1817 * @ctx: pointer to the context structure.
1818 * @device_addr:  device memory physical address.
1819 * @size: size of device memory.
1820 * @flags: DMA-BUF file/FD flags.
1821 * @dmabuf_fd: pointer to result FD that represents the dma-buf object.
1822 *
1823 * Create and export a dma-buf object for an existing memory allocation inside
1824 * the device memory, and return a FD which is associated with the dma-buf
1825 * object.
1826 *
1827 * Return: 0 on success, non-zero for failure.
1828 */
1829static int export_dmabuf_from_addr(struct hl_ctx *ctx, u64 device_addr,
1830					u64 size, int flags, int *dmabuf_fd)
1831{
1832	struct hl_dmabuf_priv *hl_dmabuf;
1833	struct hl_device *hdev = ctx->hdev;
1834	struct asic_fixed_properties *prop;
1835	u64 bar_address;
1836	int rc;
1837
1838	prop = &hdev->asic_prop;
1839
1840	if (!IS_ALIGNED(device_addr, PAGE_SIZE)) {
1841		dev_dbg(hdev->dev,
1842			"exported device memory address 0x%llx should be aligned to 0x%lx\n",
1843			device_addr, PAGE_SIZE);
1844		return -EINVAL;
1845	}
1846
1847	if (size < PAGE_SIZE) {
1848		dev_dbg(hdev->dev,
1849			"exported device memory size %llu should be equal to or greater than %lu\n",
1850			size, PAGE_SIZE);
1851		return -EINVAL;
1852	}
1853
1854	if (device_addr < prop->dram_user_base_address ||
1855				device_addr + size > prop->dram_end_address ||
1856				device_addr + size < device_addr) {
1857		dev_dbg(hdev->dev,
1858			"DRAM memory range 0x%llx (+0x%llx) is outside of DRAM boundaries\n",
1859			device_addr, size);
1860		return -EINVAL;
1861	}
1862
1863	bar_address = hdev->dram_pci_bar_start +
1864			(device_addr - prop->dram_base_address);
1865
1866	if (bar_address + size >
1867			hdev->dram_pci_bar_start + prop->dram_pci_bar_size ||
1868			bar_address + size < bar_address) {
1869		dev_dbg(hdev->dev,
1870			"DRAM memory range 0x%llx (+0x%llx) is outside of PCI BAR boundaries\n",
1871			device_addr, size);
1872		return -EINVAL;
1873	}
1874
1875	hl_dmabuf = kzalloc(sizeof(*hl_dmabuf), GFP_KERNEL);
1876	if (!hl_dmabuf)
1877		return -ENOMEM;
1878
1879	hl_dmabuf->device_address = device_addr;
1880
1881	rc = export_dmabuf_common(ctx, hl_dmabuf, size, flags, dmabuf_fd);
1882	if (rc)
1883		goto err_free_dmabuf_wrapper;
1884
1885	return 0;
1886
1887err_free_dmabuf_wrapper:
1888	kfree(hl_dmabuf);
1889	return rc;
1890}
1891
1892/**
1893 * export_dmabuf_from_handle() - export a dma-buf object for the given memory
1894 *                               handle.
1895 * @ctx: pointer to the context structure.
1896 * @handle: device memory allocation handle.
1897 * @flags: DMA-BUF file/FD flags.
1898 * @dmabuf_fd: pointer to result FD that represents the dma-buf object.
1899 *
1900 * Create and export a dma-buf object for an existing memory allocation inside
1901 * the device memory, and return a FD which is associated with the dma-buf
1902 * object.
1903 *
1904 * Return: 0 on success, non-zero for failure.
1905 */
1906static int export_dmabuf_from_handle(struct hl_ctx *ctx, u64 handle, int flags,
1907					int *dmabuf_fd)
1908{
1909	struct hl_vm_phys_pg_pack *phys_pg_pack;
1910	struct hl_dmabuf_priv *hl_dmabuf;
1911	struct hl_device *hdev = ctx->hdev;
1912	struct asic_fixed_properties *prop;
1913	struct hl_vm *vm = &hdev->vm;
1914	u64 bar_address;
1915	int rc, i;
1916
1917	prop = &hdev->asic_prop;
1918
1919	if (upper_32_bits(handle)) {
1920		dev_dbg(hdev->dev, "no match for handle 0x%llx\n", handle);
1921		return -EINVAL;
1922	}
1923
1924	spin_lock(&vm->idr_lock);
1925
1926	phys_pg_pack = idr_find(&vm->phys_pg_pack_handles, (u32) handle);
1927	if (!phys_pg_pack) {
1928		spin_unlock(&vm->idr_lock);
1929		dev_dbg(hdev->dev, "no match for handle 0x%x\n", (u32) handle);
1930		return -EINVAL;
1931	}
1932
1933	/* increment now to avoid freeing device memory while exporting */
1934	phys_pg_pack->exporting_cnt++;
1935
1936	spin_unlock(&vm->idr_lock);
1937
1938	if (phys_pg_pack->vm_type != VM_TYPE_PHYS_PACK) {
1939		dev_dbg(hdev->dev, "handle 0x%llx does not represent DRAM memory\n", handle);
1940		rc = -EINVAL;
1941		goto err_dec_exporting_cnt;
1942	}
1943
1944	for (i = 0 ; i < phys_pg_pack->npages ; i++) {
1945
1946		bar_address = hdev->dram_pci_bar_start +
1947						(phys_pg_pack->pages[i] -
1948						prop->dram_base_address);
1949
1950		if (bar_address + phys_pg_pack->page_size >
1951			hdev->dram_pci_bar_start + prop->dram_pci_bar_size ||
1952			bar_address + phys_pg_pack->page_size < bar_address) {
1953
1954			dev_dbg(hdev->dev,
1955				"DRAM memory range 0x%llx (+0x%x) is outside of PCI BAR boundaries\n",
1956				phys_pg_pack->pages[i],
1957				phys_pg_pack->page_size);
1958
1959			rc = -EINVAL;
1960			goto err_dec_exporting_cnt;
1961		}
1962	}
1963
1964	hl_dmabuf = kzalloc(sizeof(*hl_dmabuf), GFP_KERNEL);
1965	if (!hl_dmabuf) {
1966		rc = -ENOMEM;
1967		goto err_dec_exporting_cnt;
1968	}
1969
1970	hl_dmabuf->phys_pg_pack = phys_pg_pack;
1971
1972	rc = export_dmabuf_common(ctx, hl_dmabuf, phys_pg_pack->total_size,
1973				flags, dmabuf_fd);
1974	if (rc)
1975		goto err_free_dmabuf_wrapper;
1976
1977	return 0;
1978
1979err_free_dmabuf_wrapper:
1980	kfree(hl_dmabuf);
1981
1982err_dec_exporting_cnt:
1983	spin_lock(&vm->idr_lock);
1984	phys_pg_pack->exporting_cnt--;
1985	spin_unlock(&vm->idr_lock);
1986
1987	return rc;
1988}
1989
1990static int mem_ioctl_no_mmu(struct hl_fpriv *hpriv, union hl_mem_args *args)
1991{
1992	struct hl_device *hdev = hpriv->hdev;
1993	u64 block_handle, device_addr = 0;
1994	struct hl_ctx *ctx = hpriv->ctx;
1995	u32 handle = 0, block_size;
1996	int rc;
1997
1998	switch (args->in.op) {
1999	case HL_MEM_OP_ALLOC:
2000		if (args->in.alloc.mem_size == 0) {
2001			dev_err(hdev->dev, "alloc size must be larger than 0\n");
2002			rc = -EINVAL;
2003			goto out;
2004		}
2005
2006		/* Force contiguous as there are no real MMU
2007		 * translations to overcome physical memory gaps
2008		 */
2009		args->in.flags |= HL_MEM_CONTIGUOUS;
2010		rc = alloc_device_memory(ctx, &args->in, &handle);
2011
2012		memset(args, 0, sizeof(*args));
2013		args->out.handle = (__u64) handle;
2014		break;
2015
2016	case HL_MEM_OP_FREE:
2017		rc = free_device_memory(ctx, &args->in);
2018		break;
2019
2020	case HL_MEM_OP_MAP:
2021		if (args->in.flags & HL_MEM_USERPTR) {
2022			dev_err(hdev->dev, "Failed to map host memory when MMU is disabled\n");
2023			rc = -EPERM;
2024		} else {
2025			rc = get_paddr_from_handle(ctx, &args->in, &device_addr);
2026			memset(args, 0, sizeof(*args));
2027			args->out.device_virt_addr = device_addr;
2028		}
2029
2030		break;
2031
2032	case HL_MEM_OP_UNMAP:
2033		rc = 0;
2034		break;
2035
2036	case HL_MEM_OP_MAP_BLOCK:
2037		rc = map_block(hdev, args->in.map_block.block_addr, &block_handle, &block_size);
2038		args->out.block_handle = block_handle;
2039		args->out.block_size = block_size;
2040		break;
2041
2042	case HL_MEM_OP_EXPORT_DMABUF_FD:
2043		dev_err(hdev->dev, "Failed to export dma-buf object when MMU is disabled\n");
2044		rc = -EPERM;
2045		break;
2046
2047	case HL_MEM_OP_TS_ALLOC:
2048		rc = allocate_timestamps_buffers(hpriv, &args->in, &args->out.handle);
2049		break;
2050	default:
2051		dev_err(hdev->dev, "Unknown opcode for memory IOCTL\n");
2052		rc = -EINVAL;
2053		break;
2054	}
2055
2056out:
2057	return rc;
2058}
2059
2060static void ts_buff_release(struct hl_mmap_mem_buf *buf)
2061{
2062	struct hl_ts_buff *ts_buff = buf->private;
2063
2064	vfree(ts_buff->kernel_buff_address);
2065	vfree(ts_buff->user_buff_address);
2066	kfree(ts_buff);
2067}
2068
2069static int hl_ts_mmap(struct hl_mmap_mem_buf *buf, struct vm_area_struct *vma, void *args)
2070{
2071	struct hl_ts_buff *ts_buff = buf->private;
2072
2073	vma->vm_flags |= VM_DONTEXPAND | VM_DONTDUMP | VM_DONTCOPY | VM_NORESERVE;
2074	return remap_vmalloc_range(vma, ts_buff->user_buff_address, 0);
2075}
2076
2077static int hl_ts_alloc_buf(struct hl_mmap_mem_buf *buf, gfp_t gfp, void *args)
2078{
2079	struct hl_ts_buff *ts_buff = NULL;
2080	u32 size, num_elements;
2081	void *p;
2082
2083	num_elements = *(u32 *)args;
2084
2085	ts_buff = kzalloc(sizeof(*ts_buff), GFP_KERNEL);
2086	if (!ts_buff)
2087		return -ENOMEM;
2088
2089	/* Allocate the user buffer */
2090	size = num_elements * sizeof(u64);
2091	p = vmalloc_user(size);
2092	if (!p)
2093		goto free_mem;
2094
2095	ts_buff->user_buff_address = p;
2096	buf->mappable_size = size;
2097
2098	/* Allocate the internal kernel buffer */
2099	size = num_elements * sizeof(struct hl_user_pending_interrupt);
2100	p = vmalloc(size);
2101	if (!p)
2102		goto free_user_buff;
2103
2104	ts_buff->kernel_buff_address = p;
2105	ts_buff->kernel_buff_size = size;
2106
2107	buf->private = ts_buff;
2108
2109	return 0;
2110
2111free_user_buff:
2112	vfree(ts_buff->user_buff_address);
2113free_mem:
2114	kfree(ts_buff);
2115	return -ENOMEM;
2116}
2117
2118static struct hl_mmap_mem_buf_behavior hl_ts_behavior = {
2119	.topic = "TS",
2120	.mem_id = HL_MMAP_TYPE_TS_BUFF,
2121	.mmap = hl_ts_mmap,
2122	.alloc = hl_ts_alloc_buf,
2123	.release = ts_buff_release,
2124};
2125
2126/**
2127 * allocate_timestamps_buffers() - allocate timestamps buffers
2128 * This function will allocate ts buffer that will later on be mapped to the user
2129 * in order to be able to read the timestamp.
2130 * in additon it'll allocate an extra buffer for registration management.
2131 * since we cannot fail during registration for out-of-memory situation, so
2132 * we'll prepare a pool which will be used as user interrupt nodes and instead
2133 * of dynamically allocating nodes while registration we'll pick the node from
2134 * this pool. in addtion it'll add node to the mapping hash which will be used
2135 * to map user ts buffer to the internal kernel ts buffer.
2136 * @hpriv: pointer to the private data of the fd
2137 * @args: ioctl input
2138 * @handle: user timestamp buffer handle as an output
2139 */
2140static int allocate_timestamps_buffers(struct hl_fpriv *hpriv, struct hl_mem_in *args, u64 *handle)
2141{
2142	struct hl_mem_mgr *mmg = &hpriv->mem_mgr;
2143	struct hl_mmap_mem_buf *buf;
2144
2145	if (args->num_of_elements > TS_MAX_ELEMENTS_NUM) {
2146		dev_err(mmg->dev, "Num of elements exceeds Max allowed number (0x%x > 0x%x)\n",
2147				args->num_of_elements, TS_MAX_ELEMENTS_NUM);
2148		return -EINVAL;
2149	}
2150
2151	buf = hl_mmap_mem_buf_alloc(mmg, &hl_ts_behavior, GFP_KERNEL, &args->num_of_elements);
2152	if (!buf)
2153		return -ENOMEM;
2154
2155	*handle = buf->handle;
2156
2157	return 0;
2158}
2159
2160int hl_mem_ioctl(struct hl_fpriv *hpriv, void *data)
2161{
2162	enum hl_device_status status;
2163	union hl_mem_args *args = data;
2164	struct hl_device *hdev = hpriv->hdev;
2165	struct hl_ctx *ctx = hpriv->ctx;
2166	u64 block_handle, device_addr = 0;
2167	u32 handle = 0, block_size;
2168	int rc, dmabuf_fd = -EBADF;
2169
2170	if (!hl_device_operational(hdev, &status)) {
2171		dev_warn_ratelimited(hdev->dev,
2172			"Device is %s. Can't execute MEMORY IOCTL\n",
2173			hdev->status[status]);
2174		return -EBUSY;
2175	}
2176
2177	if (!hdev->mmu_enable)
2178		return mem_ioctl_no_mmu(hpriv, args);
2179
2180	switch (args->in.op) {
2181	case HL_MEM_OP_ALLOC:
2182		if (args->in.alloc.mem_size == 0) {
2183			dev_err(hdev->dev,
2184				"alloc size must be larger than 0\n");
2185			rc = -EINVAL;
2186			goto out;
2187		}
2188
2189		/* If DRAM does not support virtual memory the driver won't
2190		 * handle the allocation/freeing of that memory. However, for
2191		 * system administration/monitoring purposes, the driver will
2192		 * keep track of the amount of DRAM memory that is allocated
2193		 * and freed by the user. Because this code totally relies on
2194		 * the user's input, the driver can't ensure the validity
2195		 * of this accounting.
2196		 */
2197		if (!hdev->asic_prop.dram_supports_virtual_memory) {
2198			atomic64_add(args->in.alloc.mem_size,
2199					&ctx->dram_phys_mem);
2200			atomic64_add(args->in.alloc.mem_size,
2201					&hdev->dram_used_mem);
2202
2203			dev_dbg(hdev->dev, "DRAM alloc is not supported\n");
2204			rc = 0;
2205
2206			memset(args, 0, sizeof(*args));
2207			args->out.handle = 0;
2208			goto out;
2209		}
2210
2211		rc = alloc_device_memory(ctx, &args->in, &handle);
2212
2213		memset(args, 0, sizeof(*args));
2214		args->out.handle = (__u64) handle;
2215		break;
2216
2217	case HL_MEM_OP_FREE:
2218		/* If DRAM does not support virtual memory the driver won't
2219		 * handle the allocation/freeing of that memory. However, for
2220		 * system administration/monitoring purposes, the driver will
2221		 * keep track of the amount of DRAM memory that is allocated
2222		 * and freed by the user. Because this code totally relies on
2223		 * the user's input, the driver can't ensure the validity
2224		 * of this accounting.
2225		 */
2226		if (!hdev->asic_prop.dram_supports_virtual_memory) {
2227			atomic64_sub(args->in.alloc.mem_size,
2228					&ctx->dram_phys_mem);
2229			atomic64_sub(args->in.alloc.mem_size,
2230					&hdev->dram_used_mem);
2231
2232			dev_dbg(hdev->dev, "DRAM alloc is not supported\n");
2233			rc = 0;
2234
2235			goto out;
2236		}
2237
2238		rc = free_device_memory(ctx, &args->in);
2239		break;
2240
2241	case HL_MEM_OP_MAP:
2242		rc = map_device_va(ctx, &args->in, &device_addr);
2243
2244		memset(args, 0, sizeof(*args));
2245		args->out.device_virt_addr = device_addr;
2246		break;
2247
2248	case HL_MEM_OP_UNMAP:
2249		rc = unmap_device_va(ctx, &args->in, false);
2250		break;
2251
2252	case HL_MEM_OP_MAP_BLOCK:
2253		rc = map_block(hdev, args->in.map_block.block_addr,
2254				&block_handle, &block_size);
2255		args->out.block_handle = block_handle;
2256		args->out.block_size = block_size;
2257		break;
2258
2259	case HL_MEM_OP_EXPORT_DMABUF_FD:
2260		if (hdev->asic_prop.dram_supports_virtual_memory)
2261			rc = export_dmabuf_from_handle(ctx,
2262					args->in.export_dmabuf_fd.handle,
2263					args->in.flags,
2264					&dmabuf_fd);
2265		else
2266			rc = export_dmabuf_from_addr(ctx,
2267					args->in.export_dmabuf_fd.handle,
2268					args->in.export_dmabuf_fd.mem_size,
2269					args->in.flags,
2270					&dmabuf_fd);
2271		memset(args, 0, sizeof(*args));
2272		args->out.fd = dmabuf_fd;
2273		break;
2274
2275	case HL_MEM_OP_TS_ALLOC:
2276		rc = allocate_timestamps_buffers(hpriv, &args->in, &args->out.handle);
2277		break;
2278	default:
2279		dev_err(hdev->dev, "Unknown opcode for memory IOCTL\n");
2280		rc = -EINVAL;
2281		break;
2282	}
2283
2284out:
2285	return rc;
2286}
2287
2288static int get_user_memory(struct hl_device *hdev, u64 addr, u64 size,
2289				u32 npages, u64 start, u32 offset,
2290				struct hl_userptr *userptr)
2291{
2292	int rc;
2293
2294	if (!access_ok((void __user *) (uintptr_t) addr, size)) {
2295		dev_err(hdev->dev, "user pointer is invalid - 0x%llx\n", addr);
2296		return -EFAULT;
2297	}
2298
2299	userptr->pages = kvmalloc_array(npages, sizeof(*userptr->pages),
2300					GFP_KERNEL);
2301	if (!userptr->pages)
2302		return -ENOMEM;
2303
2304	rc = pin_user_pages_fast(start, npages,
2305				 FOLL_FORCE | FOLL_WRITE | FOLL_LONGTERM,
2306				 userptr->pages);
2307
2308	if (rc != npages) {
2309		dev_err(hdev->dev,
2310			"Failed (%d) to pin host memory with user ptr 0x%llx, size 0x%llx, npages %d\n",
2311			rc, addr, size, npages);
2312		if (rc < 0)
2313			goto destroy_pages;
2314		npages = rc;
2315		rc = -EFAULT;
2316		goto put_pages;
2317	}
2318	userptr->npages = npages;
2319
2320	rc = sg_alloc_table_from_pages(userptr->sgt,
2321				       userptr->pages,
2322				       npages, offset, size, GFP_KERNEL);
2323	if (rc < 0) {
2324		dev_err(hdev->dev, "failed to create SG table from pages\n");
2325		goto put_pages;
2326	}
2327
2328	return 0;
2329
2330put_pages:
2331	unpin_user_pages(userptr->pages, npages);
2332destroy_pages:
2333	kvfree(userptr->pages);
2334	return rc;
2335}
2336
2337/**
2338 * hl_pin_host_memory() - pins a chunk of host memory.
2339 * @hdev: pointer to the habanalabs device structure.
2340 * @addr: the host virtual address of the memory area.
2341 * @size: the size of the memory area.
2342 * @userptr: pointer to hl_userptr structure.
2343 *
2344 * This function does the following:
2345 * - Pins the physical pages.
2346 * - Create an SG list from those pages.
2347 */
2348int hl_pin_host_memory(struct hl_device *hdev, u64 addr, u64 size,
2349					struct hl_userptr *userptr)
2350{
2351	u64 start, end;
2352	u32 npages, offset;
2353	int rc;
2354
2355	if (!size) {
2356		dev_err(hdev->dev, "size to pin is invalid - %llu\n", size);
2357		return -EINVAL;
2358	}
2359
2360	/*
2361	 * If the combination of the address and size requested for this memory
2362	 * region causes an integer overflow, return error.
2363	 */
2364	if (((addr + size) < addr) ||
2365			PAGE_ALIGN(addr + size) < (addr + size)) {
2366		dev_err(hdev->dev,
2367			"user pointer 0x%llx + %llu causes integer overflow\n",
2368			addr, size);
2369		return -EINVAL;
2370	}
2371
2372	userptr->pid = current->pid;
2373	userptr->sgt = kzalloc(sizeof(*userptr->sgt), GFP_KERNEL);
2374	if (!userptr->sgt)
2375		return -ENOMEM;
2376
2377	start = addr & PAGE_MASK;
2378	offset = addr & ~PAGE_MASK;
2379	end = PAGE_ALIGN(addr + size);
2380	npages = (end - start) >> PAGE_SHIFT;
2381
2382	userptr->size = size;
2383	userptr->addr = addr;
2384	userptr->dma_mapped = false;
2385	INIT_LIST_HEAD(&userptr->job_node);
2386
2387	rc = get_user_memory(hdev, addr, size, npages, start, offset,
2388				userptr);
2389	if (rc) {
2390		dev_err(hdev->dev,
2391			"failed to get user memory for address 0x%llx\n",
2392			addr);
2393		goto free_sgt;
2394	}
2395
2396	hl_debugfs_add_userptr(hdev, userptr);
2397
2398	return 0;
2399
2400free_sgt:
2401	kfree(userptr->sgt);
2402	return rc;
2403}
2404
2405/*
2406 * hl_unpin_host_memory - unpins a chunk of host memory.
2407 * @hdev: pointer to the habanalabs device structure
2408 * @userptr: pointer to hl_userptr structure
2409 *
2410 * This function does the following:
2411 * - Unpins the physical pages related to the host memory
2412 * - Free the SG list
2413 */
2414void hl_unpin_host_memory(struct hl_device *hdev, struct hl_userptr *userptr)
2415{
2416	hl_debugfs_remove_userptr(hdev, userptr);
2417
2418	if (userptr->dma_mapped)
2419		hdev->asic_funcs->hl_dma_unmap_sgtable(hdev, userptr->sgt, userptr->dir);
2420
2421	unpin_user_pages_dirty_lock(userptr->pages, userptr->npages, true);
2422	kvfree(userptr->pages);
2423
2424	list_del(&userptr->job_node);
2425
2426	sg_free_table(userptr->sgt);
2427	kfree(userptr->sgt);
2428}
2429
2430/**
2431 * hl_userptr_delete_list() - clear userptr list.
2432 * @hdev: pointer to the habanalabs device structure.
2433 * @userptr_list: pointer to the list to clear.
2434 *
2435 * This function does the following:
2436 * - Iterates over the list and unpins the host memory and frees the userptr
2437 *   structure.
2438 */
2439void hl_userptr_delete_list(struct hl_device *hdev,
2440				struct list_head *userptr_list)
2441{
2442	struct hl_userptr *userptr, *tmp;
2443
2444	list_for_each_entry_safe(userptr, tmp, userptr_list, job_node) {
2445		hl_unpin_host_memory(hdev, userptr);
2446		kfree(userptr);
2447	}
2448
2449	INIT_LIST_HEAD(userptr_list);
2450}
2451
2452/**
2453 * hl_userptr_is_pinned() - returns whether the given userptr is pinned.
2454 * @hdev: pointer to the habanalabs device structure.
2455 * @addr: user address to check.
2456 * @size: user block size to check.
2457 * @userptr_list: pointer to the list to clear.
2458 * @userptr: pointer to userptr to check.
2459 *
2460 * This function does the following:
2461 * - Iterates over the list and checks if the given userptr is in it, means is
2462 *   pinned. If so, returns true, otherwise returns false.
2463 */
2464bool hl_userptr_is_pinned(struct hl_device *hdev, u64 addr,
2465				u32 size, struct list_head *userptr_list,
2466				struct hl_userptr **userptr)
2467{
2468	list_for_each_entry((*userptr), userptr_list, job_node) {
2469		if ((addr == (*userptr)->addr) && (size == (*userptr)->size))
2470			return true;
2471	}
2472
2473	return false;
2474}
2475
2476/**
2477 * va_range_init() - initialize virtual addresses range.
2478 * @hdev: pointer to the habanalabs device structure.
2479 * @va_ranges: pointer to va_ranges array.
2480 * @range_type: virtual address range type.
2481 * @start: range start address, inclusive.
2482 * @end: range end address, inclusive.
2483 * @page_size: page size for this va_range.
2484 *
2485 * This function does the following:
2486 * - Initializes the virtual addresses list of the given range with the given
2487 *   addresses.
2488 */
2489static int va_range_init(struct hl_device *hdev, struct hl_va_range **va_ranges,
2490				enum hl_va_range_type range_type, u64 start,
2491				u64 end, u32 page_size)
2492{
2493	struct hl_va_range *va_range = va_ranges[range_type];
2494	int rc;
2495
2496	INIT_LIST_HEAD(&va_range->list);
2497
2498	/*
2499	 * PAGE_SIZE alignment
2500	 * it is the callers responsibility to align the addresses if the
2501	 * page size is not a power of 2
2502	 */
2503
2504	if (is_power_of_2(page_size)) {
2505		if (start & (PAGE_SIZE - 1)) {
2506			start &= PAGE_MASK;
2507			start += PAGE_SIZE;
2508		}
2509
2510		/*
2511		 * The end of the range is inclusive, hence we need to align it
2512		 * to the end of the last full page in the range. For example if
2513		 * end = 0x3ff5 with page size 0x1000, we need to align it to
2514		 * 0x2fff. The remainig 0xff5 bytes do not form a full page.
2515		 */
2516		if ((end + 1) & (PAGE_SIZE - 1))
2517			end = ((end + 1) & PAGE_MASK) - 1;
2518	}
2519
2520	if (start >= end) {
2521		dev_err(hdev->dev, "too small vm range for va list\n");
2522		return -EFAULT;
2523	}
2524
2525	rc = add_va_block(hdev, va_range, start, end);
2526
2527	if (rc) {
2528		dev_err(hdev->dev, "Failed to init host va list\n");
2529		return rc;
2530	}
2531
2532	va_range->start_addr = start;
2533	va_range->end_addr = end;
2534	va_range->page_size = page_size;
2535
2536	return 0;
2537}
2538
2539/**
2540 * va_range_fini() - clear a virtual addresses range.
2541 * @hdev: pointer to the habanalabs structure.
2542 * @va_range: pointer to virtual addresses range.
2543 *
2544 * This function does the following:
2545 * - Frees the virtual addresses block list and its lock.
2546 */
2547static void va_range_fini(struct hl_device *hdev, struct hl_va_range *va_range)
2548{
2549	mutex_lock(&va_range->lock);
2550	clear_va_list_locked(hdev, &va_range->list);
2551	mutex_unlock(&va_range->lock);
2552
2553	mutex_destroy(&va_range->lock);
2554	kfree(va_range);
2555}
2556
2557/**
2558 * vm_ctx_init_with_ranges() - initialize virtual memory for context.
2559 * @ctx: pointer to the habanalabs context structure.
2560 * @host_range_start: host virtual addresses range start.
2561 * @host_range_end: host virtual addresses range end.
2562 * @host_page_size: host page size.
2563 * @host_huge_range_start: host virtual addresses range start for memory
2564 *                         allocated with huge pages.
2565 * @host_huge_range_end: host virtual addresses range end for memory allocated
2566 *                        with huge pages.
2567 * @host_huge_page_size: host huge page size.
2568 * @dram_range_start: dram virtual addresses range start.
2569 * @dram_range_end: dram virtual addresses range end.
2570 * @dram_page_size: dram page size.
2571 *
2572 * This function initializes the following:
2573 * - MMU for context.
2574 * - Virtual address to area descriptor hashtable.
2575 * - Virtual block list of available virtual memory.
2576 */
2577static int vm_ctx_init_with_ranges(struct hl_ctx *ctx,
2578					u64 host_range_start,
2579					u64 host_range_end,
2580					u32 host_page_size,
2581					u64 host_huge_range_start,
2582					u64 host_huge_range_end,
2583					u32 host_huge_page_size,
2584					u64 dram_range_start,
2585					u64 dram_range_end,
2586					u32 dram_page_size)
2587{
2588	struct hl_device *hdev = ctx->hdev;
2589	int i, rc;
2590
2591	for (i = 0 ; i < HL_VA_RANGE_TYPE_MAX ; i++) {
2592		ctx->va_range[i] =
2593			kzalloc(sizeof(struct hl_va_range), GFP_KERNEL);
2594		if (!ctx->va_range[i]) {
2595			rc = -ENOMEM;
2596			goto free_va_range;
2597		}
2598	}
2599
2600	rc = hl_mmu_ctx_init(ctx);
2601	if (rc) {
2602		dev_err(hdev->dev, "failed to init context %d\n", ctx->asid);
2603		goto free_va_range;
2604	}
2605
2606	mutex_init(&ctx->mem_hash_lock);
2607	hash_init(ctx->mem_hash);
2608
2609	mutex_init(&ctx->va_range[HL_VA_RANGE_TYPE_HOST]->lock);
2610
2611	rc = va_range_init(hdev, ctx->va_range, HL_VA_RANGE_TYPE_HOST,
2612			host_range_start, host_range_end, host_page_size);
2613	if (rc) {
2614		dev_err(hdev->dev, "failed to init host vm range\n");
2615		goto mmu_ctx_fini;
2616	}
2617
2618	if (hdev->pmmu_huge_range) {
2619		mutex_init(&ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]->lock);
2620
2621		rc = va_range_init(hdev,
2622			ctx->va_range, HL_VA_RANGE_TYPE_HOST_HUGE,
2623			host_huge_range_start, host_huge_range_end,
2624			host_huge_page_size);
2625		if (rc) {
2626			dev_err(hdev->dev,
2627				"failed to init host huge vm range\n");
2628			goto clear_host_va_range;
2629		}
2630	} else {
2631		kfree(ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]);
2632		ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE] =
2633				ctx->va_range[HL_VA_RANGE_TYPE_HOST];
2634	}
2635
2636	mutex_init(&ctx->va_range[HL_VA_RANGE_TYPE_DRAM]->lock);
2637
2638	rc = va_range_init(hdev, ctx->va_range, HL_VA_RANGE_TYPE_DRAM,
2639			dram_range_start, dram_range_end, dram_page_size);
2640	if (rc) {
2641		dev_err(hdev->dev, "failed to init dram vm range\n");
2642		goto clear_host_huge_va_range;
2643	}
2644
2645	hl_debugfs_add_ctx_mem_hash(hdev, ctx);
2646
2647	return 0;
2648
2649clear_host_huge_va_range:
2650	mutex_destroy(&ctx->va_range[HL_VA_RANGE_TYPE_DRAM]->lock);
2651
2652	if (hdev->pmmu_huge_range) {
2653		mutex_lock(&ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]->lock);
2654		clear_va_list_locked(hdev,
2655			&ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]->list);
2656		mutex_unlock(&ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]->lock);
2657	}
2658clear_host_va_range:
2659	if (hdev->pmmu_huge_range)
2660		mutex_destroy(&ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]->lock);
2661	mutex_lock(&ctx->va_range[HL_VA_RANGE_TYPE_HOST]->lock);
2662	clear_va_list_locked(hdev, &ctx->va_range[HL_VA_RANGE_TYPE_HOST]->list);
2663	mutex_unlock(&ctx->va_range[HL_VA_RANGE_TYPE_HOST]->lock);
2664mmu_ctx_fini:
2665	mutex_destroy(&ctx->va_range[HL_VA_RANGE_TYPE_HOST]->lock);
2666	mutex_destroy(&ctx->mem_hash_lock);
2667	hl_mmu_ctx_fini(ctx);
2668free_va_range:
2669	for (i = 0 ; i < HL_VA_RANGE_TYPE_MAX ; i++)
2670		kfree(ctx->va_range[i]);
2671
2672	return rc;
2673}
2674
2675int hl_vm_ctx_init(struct hl_ctx *ctx)
2676{
2677	struct asic_fixed_properties *prop = &ctx->hdev->asic_prop;
2678	u64 host_range_start, host_range_end, host_huge_range_start,
2679		host_huge_range_end, dram_range_start, dram_range_end;
2680	u32 host_page_size, host_huge_page_size, dram_page_size;
2681
2682	atomic64_set(&ctx->dram_phys_mem, 0);
2683
2684	/*
2685	 * - If MMU is enabled, init the ranges as usual.
2686	 * - If MMU is disabled, in case of host mapping, the returned address
2687	 *   is the given one.
2688	 *   In case of DRAM mapping, the returned address is the physical
2689	 *   address of the memory related to the given handle.
2690	 */
2691	if (!ctx->hdev->mmu_enable)
2692		return 0;
2693
2694	dram_range_start = prop->dmmu.start_addr;
2695	dram_range_end = prop->dmmu.end_addr - 1;
2696	dram_page_size = prop->dram_page_size ?
2697				prop->dram_page_size : prop->dmmu.page_size;
2698	host_range_start = prop->pmmu.start_addr;
2699	host_range_end = prop->pmmu.end_addr - 1;
2700	host_page_size = prop->pmmu.page_size;
2701	host_huge_range_start = prop->pmmu_huge.start_addr;
2702	host_huge_range_end = prop->pmmu_huge.end_addr - 1;
2703	host_huge_page_size = prop->pmmu_huge.page_size;
2704
2705	return vm_ctx_init_with_ranges(ctx, host_range_start, host_range_end,
2706			host_page_size, host_huge_range_start,
2707			host_huge_range_end, host_huge_page_size,
2708			dram_range_start, dram_range_end, dram_page_size);
2709}
2710
2711/**
2712 * hl_vm_ctx_fini() - virtual memory teardown of context.
2713 * @ctx: pointer to the habanalabs context structure.
2714 *
2715 * This function perform teardown the following:
2716 * - Virtual block list of available virtual memory.
2717 * - Virtual address to area descriptor hashtable.
2718 * - MMU for context.
2719 *
2720 * In addition this function does the following:
2721 * - Unmaps the existing hashtable nodes if the hashtable is not empty. The
2722 *   hashtable should be empty as no valid mappings should exist at this
2723 *   point.
2724 * - Frees any existing physical page list from the idr which relates to the
2725 *   current context asid.
2726 * - This function checks the virtual block list for correctness. At this point
2727 *   the list should contain one element which describes the whole virtual
2728 *   memory range of the context. Otherwise, a warning is printed.
2729 */
2730void hl_vm_ctx_fini(struct hl_ctx *ctx)
2731{
2732	struct hl_vm_phys_pg_pack *phys_pg_list, *tmp_phys_node;
2733	struct hl_device *hdev = ctx->hdev;
2734	struct hl_vm_hash_node *hnode;
2735	struct hl_vm *vm = &hdev->vm;
2736	struct hlist_node *tmp_node;
2737	struct list_head free_list;
2738	struct hl_mem_in args;
2739	int i;
2740
2741	if (!hdev->mmu_enable)
2742		return;
2743
2744	hl_debugfs_remove_ctx_mem_hash(hdev, ctx);
2745
2746	/*
2747	 * Clearly something went wrong on hard reset so no point in printing
2748	 * another side effect error
2749	 */
2750	if (!hdev->reset_info.hard_reset_pending && !hash_empty(ctx->mem_hash))
2751		dev_dbg(hdev->dev,
2752			"user released device without removing its memory mappings\n");
2753
2754	hash_for_each_safe(ctx->mem_hash, i, tmp_node, hnode, node) {
2755		dev_dbg(hdev->dev,
2756			"hl_mem_hash_node of vaddr 0x%llx of asid %d is still alive\n",
2757			hnode->vaddr, ctx->asid);
2758		args.unmap.device_virt_addr = hnode->vaddr;
2759		unmap_device_va(ctx, &args, true);
2760	}
2761
2762	mutex_lock(&ctx->mmu_lock);
2763
2764	/* invalidate the cache once after the unmapping loop */
2765	hl_mmu_invalidate_cache(hdev, true, MMU_OP_USERPTR);
2766	hl_mmu_invalidate_cache(hdev, true, MMU_OP_PHYS_PACK);
2767
2768	mutex_unlock(&ctx->mmu_lock);
2769
2770	INIT_LIST_HEAD(&free_list);
2771
2772	spin_lock(&vm->idr_lock);
2773	idr_for_each_entry(&vm->phys_pg_pack_handles, phys_pg_list, i)
2774		if (phys_pg_list->asid == ctx->asid) {
2775			dev_dbg(hdev->dev,
2776				"page list 0x%px of asid %d is still alive\n",
2777				phys_pg_list, ctx->asid);
2778
2779			atomic64_sub(phys_pg_list->total_size, &hdev->dram_used_mem);
2780			idr_remove(&vm->phys_pg_pack_handles, i);
2781			list_add(&phys_pg_list->node, &free_list);
2782		}
2783	spin_unlock(&vm->idr_lock);
2784
2785	list_for_each_entry_safe(phys_pg_list, tmp_phys_node, &free_list, node)
2786		free_phys_pg_pack(hdev, phys_pg_list);
2787
2788	va_range_fini(hdev, ctx->va_range[HL_VA_RANGE_TYPE_DRAM]);
2789	va_range_fini(hdev, ctx->va_range[HL_VA_RANGE_TYPE_HOST]);
2790
2791	if (hdev->pmmu_huge_range)
2792		va_range_fini(hdev, ctx->va_range[HL_VA_RANGE_TYPE_HOST_HUGE]);
2793
2794	mutex_destroy(&ctx->mem_hash_lock);
2795	hl_mmu_ctx_fini(ctx);
2796
2797	/* In this case we need to clear the global accounting of DRAM usage
2798	 * because the user notifies us on allocations. If the user is no more,
2799	 * all DRAM is available
2800	 */
2801	if (ctx->asid != HL_KERNEL_ASID_ID &&
2802			!hdev->asic_prop.dram_supports_virtual_memory)
2803		atomic64_set(&hdev->dram_used_mem, 0);
2804}
2805
2806/**
2807 * hl_vm_init() - initialize virtual memory module.
2808 * @hdev: pointer to the habanalabs device structure.
2809 *
2810 * This function initializes the following:
2811 * - MMU module.
2812 * - DRAM physical pages pool of 2MB.
2813 * - Idr for device memory allocation handles.
2814 */
2815int hl_vm_init(struct hl_device *hdev)
2816{
2817	struct asic_fixed_properties *prop = &hdev->asic_prop;
2818	struct hl_vm *vm = &hdev->vm;
2819	int rc;
2820
2821	if (is_power_of_2(prop->dram_page_size))
2822		vm->dram_pg_pool =
2823			gen_pool_create(__ffs(prop->dram_page_size), -1);
2824	else
2825		vm->dram_pg_pool =
2826			gen_pool_create(__ffs(DRAM_POOL_PAGE_SIZE), -1);
2827
2828	if (!vm->dram_pg_pool) {
2829		dev_err(hdev->dev, "Failed to create dram page pool\n");
2830		return -ENOMEM;
2831	}
2832
2833	kref_init(&vm->dram_pg_pool_refcount);
2834
2835	rc = gen_pool_add(vm->dram_pg_pool, prop->dram_user_base_address,
2836			prop->dram_end_address - prop->dram_user_base_address,
2837			-1);
2838
2839	if (rc) {
2840		dev_err(hdev->dev,
2841			"Failed to add memory to dram page pool %d\n", rc);
2842		goto pool_add_err;
2843	}
2844
2845	spin_lock_init(&vm->idr_lock);
2846	idr_init(&vm->phys_pg_pack_handles);
2847
2848	atomic64_set(&hdev->dram_used_mem, 0);
2849
2850	vm->init_done = true;
2851
2852	return 0;
2853
2854pool_add_err:
2855	gen_pool_destroy(vm->dram_pg_pool);
2856
2857	return rc;
2858}
2859
2860/**
2861 * hl_vm_fini() - virtual memory module teardown.
2862 * @hdev: pointer to the habanalabs device structure.
2863 *
2864 * This function perform teardown to the following:
2865 * - Idr for device memory allocation handles.
2866 * - DRAM physical pages pool of 2MB.
2867 * - MMU module.
2868 */
2869void hl_vm_fini(struct hl_device *hdev)
2870{
2871	struct hl_vm *vm = &hdev->vm;
2872
2873	if (!vm->init_done)
2874		return;
2875
2876	/*
2877	 * At this point all the contexts should be freed and hence no DRAM
2878	 * memory should be in use. Hence the DRAM pool should be freed here.
2879	 */
2880	if (kref_put(&vm->dram_pg_pool_refcount, dram_pg_pool_do_release) != 1)
2881		dev_warn(hdev->dev, "dram_pg_pool was not destroyed on %s\n",
2882				__func__);
2883
2884	vm->init_done = false;
2885}
2886
2887/**
2888 * hl_hw_block_mem_init() - HW block memory initialization.
2889 * @ctx: pointer to the habanalabs context structure.
2890 *
2891 * This function initializes the HW block virtual mapped addresses list and
2892 * it's lock.
2893 */
2894void hl_hw_block_mem_init(struct hl_ctx *ctx)
2895{
2896	mutex_init(&ctx->hw_block_list_lock);
2897	INIT_LIST_HEAD(&ctx->hw_block_mem_list);
2898}
2899
2900/**
2901 * hl_hw_block_mem_fini() - HW block memory teardown.
2902 * @ctx: pointer to the habanalabs context structure.
2903 *
2904 * This function clears the HW block virtual mapped addresses list and destroys
2905 * it's lock.
2906 */
2907void hl_hw_block_mem_fini(struct hl_ctx *ctx)
2908{
2909	struct hl_vm_hw_block_list_node *lnode, *tmp;
2910
2911	if (!list_empty(&ctx->hw_block_mem_list))
2912		dev_crit(ctx->hdev->dev, "HW block mem list isn't empty\n");
2913
2914	list_for_each_entry_safe(lnode, tmp, &ctx->hw_block_mem_list, node) {
2915		list_del(&lnode->node);
2916		kfree(lnode);
2917	}
2918
2919	mutex_destroy(&ctx->hw_block_list_lock);
2920}
2921