habanalabs: Fix virtual address access via debugfs for 2MB pages

The debugfs interface for accessing DRAM virtual addresses currently
uses the 12 LSBs of a virtual address as an offset.
However, it should use the 20 LSBs in case the device MMU page size is
2MB instead of 4KB.
This patch fixes the offset calculation to be based on the page size.
Signed-off-by: Tomer Tayar <ttayar@habana.ai>
Reviewed-by: Oded Gabbay <oded.gabbay@gmail.com>
Signed-off-by: Oded Gabbay <oded.gabbay@gmail.com>

drivers/misc/habanalabs/debugfs.c

@@ -500,6 +500,7 @@ static int device_va_to_pa(struct hl_device *hdev, u64 virt_addr,
 {
 	struct hl_ctx *ctx = hdev->user_ctx;
 	u64 hop_addr, hop_pte_addr, hop_pte;
+	u64 offset_mask = HOP4_MASK | OFFSET_MASK;
 	int rc = 0;
 
 	if (!ctx) {
@@ -542,12 +543,14 @@ static int device_va_to_pa(struct hl_device *hdev, u64 virt_addr,
 			goto not_mapped;
 		hop_pte_addr = get_hop4_pte_addr(ctx, hop_addr, virt_addr);
 		hop_pte = hdev->asic_funcs->read_pte(hdev, hop_pte_addr);
+
+		offset_mask = OFFSET_MASK;
 	}
 
 	if (!(hop_pte & PAGE_PRESENT_MASK))
 		goto not_mapped;
 
-	*phys_addr = (hop_pte & PTE_PHYS_ADDR_MASK) | (virt_addr & OFFSET_MASK);
+	*phys_addr = (hop_pte & ~offset_mask) | (virt_addr & offset_mask);
 
 	goto out;