1 files changed, 25 insertions, 0 deletions

diff --git a/runlist.c b/runlist.c
index 8dfa1c7..03528af 100644
--- a/runlist.c
+++ b/runlist.c
@@ -109,3 +109,28 @@ int get_runlist_iter(struct runlist_iter *rl_iter) {
         printk(KERN_INFO "[nvdebug] tsg_length: %d\n", head.tsg_length);
         printk(KERN_INFO "[nvdebug] tsgid: %d\n", head.tsgid); */
 }
+
+int preempt_tsg(uint32_t tsg_id) {
+        struct gk20a *g = get_live_gk20a();
+        runlist_info_t rl_info;
+        pfifo_preempt_t pfifo_preempt;
+        runlist_disable_t rl_disable;
+        if (!g)
+                return -EIO;
+        rl_info.raw = nvdebug_readl(g, NV_PFIFO_RUNLIST);
+        pfifo_preempt.id = tsg_id;
+        pfifo_preempt.is_pending = 0;
+        pfifo_preempt.type = PREEMPT_TYPE_TSG;
+        // There may be a bug (?) that requires us to disable scheduling before preempting
+        rl_disable.raw = nvdebug_readl(g, NV_PFIFO_SCHED_DISABLE);
+        rl_disable.raw |= BIT(rl_info.id);  // Disable runlist rl_info.id
+        nvdebug_writel(g, NV_PFIFO_SCHED_DISABLE, rl_disable.raw);
+        // Actually trigger the preemption
+        nvdebug_writel(g, NV_PFIFO_PREEMPT, pfifo_preempt.raw);
+        // Renable scheduling
+        rl_disable.raw &= ~BIT(rl_info.id);  // Enable runlist rl_info.id
+        nvdebug_writel(g, NV_PFIFO_SCHED_DISABLE, rl_disable.raw);
+
+        printk(KERN_INFO "[nvdebug] TSG %d preempted (runlist %d)\n", tsg_id, rl_info.id);
+        return 0;
+}