habanalabs: fix debugfs address translation

when user uses virtual addresses to access dram through debugfs,
driver translate this address to physical and use it
for the access through the pcie bar.
in case dram page size is different than the dmmu
page size, we need to have special treatment
for adding the page offset to the actual address, which
is to use the dram page size mask to fetch the page offset
from the virtual address, instead of the dmmu last hop shift.

Signed-off-by: farah kassabri <fkassabri@habana.ai>
Reviewed-by: Oded Gabbay <ogabbay@kernel.org>
Signed-off-by: Oded Gabbay <ogabbay@kernel.org>

Author: Farah-kassabri

drivers/misc/habanalabs/common/mmu/mmu.c

@@ -499,18 +499,32 @@ static void hl_mmu_pa_page_with_offset(struct hl_ctx *ctx, u64 virt_addr,
 	else /* HL_VA_RANGE_TYPE_DRAM */
 		p = &prop->dmmu;
 
-	/*
-	 * find the correct hop shift field in hl_mmu_properties structure
-	 * in order to determine the right maks for the page offset.
-	 */
-	hop0_shift_off = offsetof(struct hl_mmu_properties, hop0_shift);
-	p = (char *)p + hop0_shift_off;
-	p = (char *)p + ((hops->used_hops - 1) * sizeof(u64));
-	hop_shift = *(u64 *)p;
-	offset_mask = (1ull << hop_shift) - 1;
-	addr_mask = ~(offset_mask);
-	*phys_addr = (tmp_phys_addr & addr_mask) |
-					(virt_addr & offset_mask);
+	if ((hops->range_type == HL_VA_RANGE_TYPE_DRAM) &&
+			!is_power_of_2(prop->dram_page_size)) {
+		u32 bit;
+		u64 page_offset_mask;
+		u64 phys_addr_mask;
+
+		bit = __ffs64((u64)prop->dram_page_size);
+		page_offset_mask = ((1ull << bit) - 1);
+		phys_addr_mask = ~page_offset_mask;
+		*phys_addr = (tmp_phys_addr & phys_addr_mask) |
+				(virt_addr & page_offset_mask);
+	} else {
+		/*
+		 * find the correct hop shift field in hl_mmu_properties
+		 * structure in order to determine the right masks
+		 * for the page offset.
+		 */
+		hop0_shift_off = offsetof(struct hl_mmu_properties, hop0_shift);
+		p = (char *)p + hop0_shift_off;
+		p = (char *)p + ((hops->used_hops - 1) * sizeof(u64));
+		hop_shift = *(u64 *)p;
+		offset_mask = (1ull << hop_shift) - 1;
+		addr_mask = ~(offset_mask);
+		*phys_addr = (tmp_phys_addr & addr_mask) |
+				(virt_addr & offset_mask);
+	}
 }
 
 int hl_mmu_va_to_pa(struct hl_ctx *ctx, u64 virt_addr, u64 *phys_addr)