diff --git a/drivers/avr/ws2812.c b/drivers/avr/ws2812.c
index 82d985c20a..5c3d72dcb5 100644
--- a/drivers/avr/ws2812.c
+++ b/drivers/avr/ws2812.c
@@ -20,12 +20,13 @@
  * You should have received a copy of the GNU General Public License
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
-
 #include "ws2812.h"
 #include <avr/interrupt.h>
 #include <avr/io.h>
 #include <util/delay.h>
 
+#define pinmask(pin) (_BV((pin)&0xF))
+
 /*
  * Forward declare internal functions
  *
@@ -33,20 +34,21 @@
  * The length is the number of bytes to send - three per LED.
  */
 
-void ws2812_sendarray(uint8_t *array, uint16_t length);
-void ws2812_sendarray_mask(uint8_t *array, uint16_t length, uint8_t pinmask);
+static inline void ws2812_sendarray_mask(uint8_t *data, uint16_t datlen, uint8_t masklo, uint8_t maskhi);
 
 // Setleds for standard RGB
-void inline ws2812_setleds(LED_TYPE *ledarray, uint16_t leds) {
-    // ws2812_setleds_pin(ledarray,leds, _BV(ws2812_pin));
-    ws2812_setleds_pin(ledarray, leds, _BV(RGB_DI_PIN & 0xF));
+void inline ws2812_setleds(LED_TYPE *ledarray, uint16_t number_of_leds) {
+    // wrap up usage of RGB_DI_PIN
+    ws2812_setleds_pin(ledarray, number_of_leds, RGB_DI_PIN);
 }
 
-void inline ws2812_setleds_pin(LED_TYPE *ledarray, uint16_t leds, uint8_t pinmask) {
-    // new universal format (DDR)
-    _SFR_IO8((RGB_DI_PIN >> 4) + 1) |= pinmask;
+void ws2812_setleds_pin(LED_TYPE *ledarray, uint16_t number_of_leds, uint8_t pin) {
+    DDRx_ADDRESS(RGB_DI_PIN) |= pinmask(pin);
 
-    ws2812_sendarray_mask((uint8_t *)ledarray, leds * sizeof(LED_TYPE), pinmask);
+    uint8_t masklo = ~(pinmask(pin)) & PORTx_ADDRESS(pin);
+    uint8_t maskhi = pinmask(pin) | PORTx_ADDRESS(pin);
+
+    ws2812_sendarray_mask((uint8_t *)ledarray, number_of_leds * sizeof(LED_TYPE), masklo, maskhi);
 
 #ifdef RGBW
     _delay_us(80);
@@ -55,8 +57,6 @@ void inline ws2812_setleds_pin(LED_TYPE *ledarray, uint16_t leds, uint8_t pinmas
 #endif
 }
 
-void ws2812_sendarray(uint8_t *data, uint16_t datlen) { ws2812_sendarray_mask(data, datlen, _BV(RGB_DI_PIN & 0xF)); }
-
 /*
   This routine writes an array of bytes with RGB values to the Dataout pin
   using the fast 800kHz clockless WS2811/2812 protocol.
@@ -118,14 +118,9 @@ void ws2812_sendarray(uint8_t *data, uint16_t datlen) { ws2812_sendarray_mask(da
 #define w_nop8 w_nop4 w_nop4
 #define w_nop16 w_nop8 w_nop8
 
-void inline ws2812_sendarray_mask(uint8_t *data, uint16_t datlen, uint8_t maskhi) {
-    uint8_t curbyte, ctr, masklo;
-    uint8_t sreg_prev;
+static inline void ws2812_sendarray_mask(uint8_t *data, uint16_t datlen, uint8_t masklo, uint8_t maskhi) {
+    uint8_t curbyte, ctr, sreg_prev;
 
-    // masklo  =~maskhi&ws2812_PORTREG;
-    // maskhi |=        ws2812_PORTREG;
-    masklo = ~maskhi & _SFR_IO8((RGB_DI_PIN >> 4) + 2);
-    maskhi |= _SFR_IO8((RGB_DI_PIN >> 4) + 2);
     sreg_prev = SREG;
     cli();
 
@@ -188,7 +183,7 @@ void inline ws2812_sendarray_mask(uint8_t *data, uint16_t datlen, uint8_t maskhi
                      "       dec   %0    \n\t"  //  '1' [+2] '0' [+2]
                      "       brne  loop%=\n\t"  //  '1' [+3] '0' [+4]
                      : "=&d"(ctr)
-                     : "r"(curbyte), "I"(_SFR_IO_ADDR(_SFR_IO8((RGB_DI_PIN >> 4) + 2))), "r"(maskhi), "r"(masklo));
+                     : "r"(curbyte), "I"(_SFR_IO_ADDR(PORTx_ADDRESS(RGB_DI_PIN))), "r"(maskhi), "r"(masklo));
     }
 
     SREG = sreg_prev;
diff --git a/drivers/avr/ws2812.h b/drivers/avr/ws2812.h
index b869fb28c8..88eb081894 100644
--- a/drivers/avr/ws2812.h
+++ b/drivers/avr/ws2812.h
@@ -29,7 +29,7 @@
  * Input:
  *         ledarray:           An array of GRB data describing the LED colors
  *         number_of_leds:     The number of LEDs to write
- *         pinmask (optional): Bitmask describing the output bin. e.g. _BV(PB0)
+ *         pin (optional):     A pin_t definition for the line to drive
  *
  * The functions will perform the following actions:
  *         - Set the data-out pin as output
@@ -37,4 +37,4 @@
  *         - Wait 50us to reset the LEDs
  */
 void ws2812_setleds(LED_TYPE *ledarray, uint16_t number_of_leds);
-void ws2812_setleds_pin(LED_TYPE *ledarray, uint16_t number_of_leds, uint8_t pinmask);
+void ws2812_setleds_pin(LED_TYPE *ledarray, uint16_t number_of_leds, uint8_t pin);
diff --git a/keyboards/ergodox_ez/led_i2c.c b/keyboards/ergodox_ez/led_i2c.c
index 4a7a02f468..7c1ccdec5a 100644
--- a/keyboards/ergodox_ez/led_i2c.c
+++ b/keyboards/ergodox_ez/led_i2c.c
@@ -27,19 +27,6 @@
 
 extern rgblight_config_t rgblight_config;
 
-/*
- * Forward declare internal functions
- *
- * The functions take a byte-array and send to the data output as WS2812 bitstream.
- * The length is the number of bytes to send - three per LED.
- */
-
-void ws2812_sendarray(uint8_t *array, uint16_t length);
-void ws2812_sendarray_mask(uint8_t *array, uint16_t length, uint8_t pinmask);
-
-
-
-
 void rgblight_set(void) {
     if (!rgblight_config.enable) {
         for (uint8_t i = 0; i < RGBLED_NUM; i++) {