This is required to allow the IOMMU driver to correctly flush its own
TLB.

Signed-off-by: Alexandre Ghiti <alexghiti@rivosinc.com>
Reviewed-by: Clément Léger <cleger@rivosinc.com>
---

Changes in v3:
- Rebase on top of 6.13-rc3

Changes in v2:
- Rebase on top of 6.9-rc1

 arch/riscv/mm/tlbflush.c | 39 ++++++++++++++++++++++++---------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/arch/riscv/mm/tlbflush.c b/arch/riscv/mm/tlbflush.c
index 9b6e86ce3867..f5f8bc6d1c61 100644
--- a/arch/riscv/mm/tlbflush.c
+++ b/arch/riscv/mm/tlbflush.c
@@ -4,6 +4,7 @@
 #include <linux/smp.h>
 #include <linux/sched.h>
 #include <linux/hugetlb.h>
+#include <linux/mmu_notifier.h>
 #include <asm/sbi.h>
 #include <asm/mmu_context.h>
 
@@ -78,10 +79,19 @@ static void __ipi_flush_tlb_range_asid(void *info)
 	local_flush_tlb_range_asid(d->start, d->size, d->stride, d->asid);
 }
 
-static void __flush_tlb_range(const struct cpumask *cmask, unsigned long asid,
+static inline unsigned long get_mm_asid(struct mm_struct *mm)
+{
+	return mm ? cntx2asid(atomic_long_read(&mm->context.id)) : FLUSH_TLB_NO_ASID;
+}
+
+static void __flush_tlb_range(struct mm_struct *mm,
+			      const struct cpumask *cmask,
 			      unsigned long start, unsigned long size,
 			      unsigned long stride)
 {
+	struct flush_tlb_range_data ftd;
+	unsigned long asid = get_mm_asid(mm);
+	bool broadcast;
 	unsigned int cpu;
 
 	if (cpumask_empty(cmask))
@@ -105,30 +115,26 @@ static void __flush_tlb_range(const struct cpumask *cmask, unsigned long asid,
 	}
 
 	put_cpu();
-}
 
-static inline unsigned long get_mm_asid(struct mm_struct *mm)
-{
-	return cntx2asid(atomic_long_read(&mm->context.id));
+	if (mm)
+		mmu_notifier_arch_invalidate_secondary_tlbs(mm, start, start + size);
 }
 
 void flush_tlb_mm(struct mm_struct *mm)
 {
-	__flush_tlb_range(mm_cpumask(mm), get_mm_asid(mm),
-			  0, FLUSH_TLB_MAX_SIZE, PAGE_SIZE);
+	__flush_tlb_range(mm, mm_cpumask(mm), 0, FLUSH_TLB_MAX_SIZE, PAGE_SIZE);
 }
 
 void flush_tlb_mm_range(struct mm_struct *mm,
 			unsigned long start, unsigned long end,
 			unsigned int page_size)
 {
-	__flush_tlb_range(mm_cpumask(mm), get_mm_asid(mm),
-			  start, end - start, page_size);
+	__flush_tlb_range(mm, mm_cpumask(mm), start, end - start, page_size);
 }
 
 void flush_tlb_page(struct vm_area_struct *vma, unsigned long addr)
 {
-	__flush_tlb_range(mm_cpumask(vma->vm_mm), get_mm_asid(vma->vm_mm),
+	__flush_tlb_range(vma->vm_mm, mm_cpumask(vma->vm_mm),
 			  addr, PAGE_SIZE, PAGE_SIZE);
 }
 
@@ -161,13 +167,13 @@ void flush_tlb_range(struct vm_area_struct *vma, unsigned long start,
 		}
 	}
 
-	__flush_tlb_range(mm_cpumask(vma->vm_mm), get_mm_asid(vma->vm_mm),
+	__flush_tlb_range(vma->vm_mm, mm_cpumask(vma->vm_mm),
 			  start, end - start, stride_size);
 }
 
 void flush_tlb_kernel_range(unsigned long start, unsigned long end)
 {
-	__flush_tlb_range(cpu_online_mask, FLUSH_TLB_NO_ASID,
+	__flush_tlb_range(NULL, cpu_online_mask,
 			  start, end - start, PAGE_SIZE);
 }
 
@@ -175,7 +181,7 @@ void flush_tlb_kernel_range(unsigned long start, unsigned long end)
 void flush_pmd_tlb_range(struct vm_area_struct *vma, unsigned long start,
 			unsigned long end)
 {
-	__flush_tlb_range(mm_cpumask(vma->vm_mm), get_mm_asid(vma->vm_mm),
+	__flush_tlb_range(vma->vm_mm, mm_cpumask(vma->vm_mm),
 			  start, end - start, PMD_SIZE);
 }
 #endif
@@ -189,7 +195,10 @@ void arch_tlbbatch_add_pending(struct arch_tlbflush_unmap_batch *batch,
 			       struct mm_struct *mm,
 			       unsigned long uaddr)
 {
+	unsigned long start = uaddr & PAGE_MASK;
+
 	cpumask_or(&batch->cpumask, &batch->cpumask, mm_cpumask(mm));
+	mmu_notifier_arch_invalidate_secondary_tlbs(mm, start, start + PAGE_SIZE);
 }
 
 void arch_flush_tlb_batched_pending(struct mm_struct *mm)
@@ -199,7 +208,7 @@ void arch_flush_tlb_batched_pending(struct mm_struct *mm)
 
 void arch_tlbbatch_flush(struct arch_tlbflush_unmap_batch *batch)
 {
-	__flush_tlb_range(&batch->cpumask, FLUSH_TLB_NO_ASID, 0,
-			  FLUSH_TLB_MAX_SIZE, PAGE_SIZE);
+	__flush_tlb_range(NULL, &batch->cpumask,
+			  0, FLUSH_TLB_MAX_SIZE, PAGE_SIZE);
 	cpumask_clear(&batch->cpumask);
 }
-- 
2.39.2

Hi Alexandre,

kernel test robot noticed the following build warnings:

[auto build test WARNING on linus/master]
[also build test WARNING on v6.13-rc6 next-20250109]
[If your patch is applied to the wrong git tree, kindly drop us a note.
And when submitting patch, we suggest to use '--base' as documented in
https://git-scm.com/docs/git-format-patch#_base_tree_information]

url:    https://github.com/intel-lab-lkp/linux/commits/Alexandre-Ghiti/riscv-Call-secondary-mmu-notifier-when-flushing-the-tlb/20250108-235624
base:   linus/master
patch link:    https://lore.kernel.org/r/20250108154907.132863-1-alexghiti%40rivosinc.com
patch subject: [PATCH v3] riscv: Call secondary mmu notifier when flushing the tlb
config: riscv-allnoconfig (https://download.01.org/0day-ci/archive/20250109/202501092121.lZ2DY62V-lkp@intel.com/config)
compiler: riscv64-linux-gcc (GCC) 14.2.0
reproduce (this is a W=1 build): (https://download.01.org/0day-ci/archive/20250109/202501092121.lZ2DY62V-lkp@intel.com/reproduce)

If you fix the issue in a separate patch/commit (i.e. not just a new version of
the same patch/commit), kindly add following tags
| Reported-by: kernel test robot <lkp@intel.com>
| Closes: https://lore.kernel.org/oe-kbuild-all/202501092121.lZ2DY62V-lkp@intel.com/

All warnings (new ones prefixed by >>):

   arch/riscv/mm/tlbflush.c: In function '__flush_tlb_range':
>> arch/riscv/mm/tlbflush.c:94:14: warning: unused variable 'broadcast' [-Wunused-variable]
      94 |         bool broadcast;
         |              ^~~~~~~~~
>> arch/riscv/mm/tlbflush.c:92:37: warning: unused variable 'ftd' [-Wunused-variable]
      92 |         struct flush_tlb_range_data ftd;
         |                                     ^~~


vim +/broadcast +94 arch/riscv/mm/tlbflush.c

    86	
    87	static void __flush_tlb_range(struct mm_struct *mm,
    88				      const struct cpumask *cmask,
    89				      unsigned long start, unsigned long size,
    90				      unsigned long stride)
    91	{
  > 92		struct flush_tlb_range_data ftd;
    93		unsigned long asid = get_mm_asid(mm);
  > 94		bool broadcast;
    95		unsigned int cpu;
    96	
    97		if (cpumask_empty(cmask))
    98			return;
    99	
   100		cpu = get_cpu();
   101	
   102		/* Check if the TLB flush needs to be sent to other CPUs. */
   103		if (cpumask_any_but(cmask, cpu) >= nr_cpu_ids) {
   104			local_flush_tlb_range_asid(start, size, stride, asid);
   105		} else if (riscv_use_sbi_for_rfence()) {
   106			sbi_remote_sfence_vma_asid(cmask, start, size, asid);
   107		} else {
   108			struct flush_tlb_range_data ftd;
   109	
   110			ftd.asid = asid;
   111			ftd.start = start;
   112			ftd.size = size;
   113			ftd.stride = stride;
   114			on_each_cpu_mask(cmask, __ipi_flush_tlb_range_asid, &ftd, 1);
   115		}
   116	
   117		put_cpu();
   118	
   119		if (mm)
   120			mmu_notifier_arch_invalidate_secondary_tlbs(mm, start, start + size);
   121	}
   122	

-- 
0-DAY CI Kernel Test Service
https://github.com/intel/lkp-tests/wiki