off_cpu_bpf__destroy(skel);
 }
 
+/* v5.18 kernel added prev_state arg, so it needs to check the signature */
+static void check_sched_switch_args(void)
+{
+       const struct btf *btf = bpf_object__btf(skel->obj);
+       const struct btf_type *t1, *t2, *t3;
+       u32 type_id;
+
+       type_id = btf__find_by_name_kind(btf, "bpf_trace_sched_switch",
+                                        BTF_KIND_TYPEDEF);
+       if ((s32)type_id < 0)
+               return;
+
+       t1 = btf__type_by_id(btf, type_id);
+       if (t1 == NULL)
+               return;
+
+       t2 = btf__type_by_id(btf, t1->type);
+       if (t2 == NULL || !btf_is_ptr(t2))
+               return;
+
+       t3 = btf__type_by_id(btf, t2->type);
+       if (t3 && btf_is_func_proto(t3) && btf_vlen(t3) == 4) {
+               /* new format: pass prev_state as 4th arg */
+               skel->rodata->has_prev_state = true;
+       }
+}
+
 int off_cpu_prepare(struct evlist *evlist, struct target *target)
 {
        int err, fd, i;
        }
 
        set_max_rlimit();
+       check_sched_switch_args();
 
        err = off_cpu_bpf__load(skel);
        if (err) {
 
 int has_cpu = 0;
 int has_task = 0;
 
+const volatile bool has_prev_state = false;
+
 /*
  * Old kernel used to call it task_struct->state and now it's '__state'.
  * Use BPF CO-RE "ignored suffix rule" to deal with it like below:
        return 1;
 }
 
-SEC("tp_btf/sched_switch")
-int on_switch(u64 *ctx)
+static int off_cpu_stat(u64 *ctx, struct task_struct *prev,
+                       struct task_struct *next, int state)
 {
        __u64 ts;
-       int state;
        __u32 stack_id;
-       struct task_struct *prev, *next;
        struct tstamp_data *pelem;
 
-       if (!enabled)
-               return 0;
-
-       prev = (struct task_struct *)ctx[1];
-       next = (struct task_struct *)ctx[2];
-       state = get_task_state(prev);
-
        ts = bpf_ktime_get_ns();
 
        if (!can_record(prev, state))
        return 0;
 }
 
+SEC("tp_btf/sched_switch")
+int on_switch(u64 *ctx)
+{
+       struct task_struct *prev, *next;
+       int prev_state;
+
+       if (!enabled)
+               return 0;
+
+       prev = (struct task_struct *)ctx[1];
+       next = (struct task_struct *)ctx[2];
+
+       if (has_prev_state)
+               prev_state = (int)ctx[3];
+       else
+               prev_state = get_task_state(prev);
+
+       return off_cpu_stat(ctx, prev, next, prev_state);
+}
+
 char LICENSE[] SEC("license") = "Dual BSD/GPL";