diff --git a/drivers/clocksource/timer-atmel-pit.c b/drivers/clocksource/timer-atmel-pit.c
index c0304ff..5875322 100644
--- a/drivers/clocksource/timer-atmel-pit.c
+++ b/drivers/clocksource/timer-atmel-pit.c
@@ -90,33 +90,27 @@
 	return elapsed;
 }
 
-/*
- * Clockevent device:  interrupts every 1/HZ (== pit_cycles * MCK/16)
- */
-static void
-pit_clkevt_mode(enum clock_event_mode mode, struct clock_event_device *dev)
+static int pit_clkevt_shutdown(struct clock_event_device *dev)
 {
 	struct pit_data *data = clkevt_to_pit_data(dev);
 
-	switch (mode) {
-	case CLOCK_EVT_MODE_PERIODIC:
-		/* update clocksource counter */
-		data->cnt += data->cycle * PIT_PICNT(pit_read(data->base, AT91_PIT_PIVR));
-		pit_write(data->base, AT91_PIT_MR,
-			  (data->cycle - 1) | AT91_PIT_PITEN | AT91_PIT_PITIEN);
-		break;
-	case CLOCK_EVT_MODE_ONESHOT:
-		BUG();
-		/* FALLTHROUGH */
-	case CLOCK_EVT_MODE_SHUTDOWN:
-	case CLOCK_EVT_MODE_UNUSED:
-		/* disable irq, leaving the clocksource active */
-		pit_write(data->base, AT91_PIT_MR,
-			  (data->cycle - 1) | AT91_PIT_PITEN);
-		break;
-	case CLOCK_EVT_MODE_RESUME:
-		break;
-	}
+	/* disable irq, leaving the clocksource active */
+	pit_write(data->base, AT91_PIT_MR, (data->cycle - 1) | AT91_PIT_PITEN);
+	return 0;
+}
+
+/*
+ * Clockevent device:  interrupts every 1/HZ (== pit_cycles * MCK/16)
+ */
+static int pit_clkevt_set_periodic(struct clock_event_device *dev)
+{
+	struct pit_data *data = clkevt_to_pit_data(dev);
+
+	/* update clocksource counter */
+	data->cnt += data->cycle * PIT_PICNT(pit_read(data->base, AT91_PIT_PIVR));
+	pit_write(data->base, AT91_PIT_MR,
+		  (data->cycle - 1) | AT91_PIT_PITEN | AT91_PIT_PITIEN);
+	return 0;
 }
 
 static void at91sam926x_pit_suspend(struct clock_event_device *cedev)
@@ -162,7 +156,7 @@
 	WARN_ON_ONCE(!irqs_disabled());
 
 	/* The PIT interrupt may be disabled, and is shared */
-	if ((data->clkevt.mode == CLOCK_EVT_MODE_PERIODIC) &&
+	if (clockevent_state_periodic(&data->clkevt) &&
 	    (pit_read(data->base, AT91_PIT_SR) & AT91_PIT_PITS)) {
 		unsigned nr_ticks;
 
@@ -227,7 +221,8 @@
 	data->clkevt.rating = 100;
 	data->clkevt.cpumask = cpumask_of(0);
 
-	data->clkevt.set_mode = pit_clkevt_mode;
+	data->clkevt.set_state_shutdown = pit_clkevt_shutdown;
+	data->clkevt.set_state_periodic = pit_clkevt_set_periodic;
 	data->clkevt.resume = at91sam926x_pit_resume;
 	data->clkevt.suspend = at91sam926x_pit_suspend;
 	clockevents_register_device(&data->clkevt);