SSD1306 display goes out-of-sync only if ATmega328 is reprogrammed

Question

I have ATmega328 directly connected to the SSD1306 Display through I2C. Whenever I power up the ATmega, everything works great and the display works as expected.

The problem occurs only when I upload a new program to the ATmega while the ATmega is already running. Then the display goes out of sync ether horizontally or vertically and starts at a random offset.

This is the expected result, I can power it up 50x in a row and it still is in sync (rendering starts from 1st pixel of x coordinate):

This is how it acts when I re-program it, it goes out-of-sync ether vertically or horizontally (rendering starts from random location) (if I power it off-on it will work fine again):

I'm not sure how to solve this or if this is a software or hardware issue. Any hints would be much appreciated. I will post the code just in case:

twi.c

void TWIInit()
{
    TWSR = 0x00;          // prescaler
    TWBR = 0x0C;          // divider
    TWCR = ( 1 << TWEN ); // enable TWI
}


void TWIStart( void )
{
    TWCR = ( 1 << TWINT ) | ( 1 << TWSTA ) | ( 1 << TWEN ) | ( 1 << TWEA );

    while ( ( TWCR & ( 1 << TWINT ) ) == 0 );
}


void TWIStop( void )
{
    TWCR = ( 1 << TWINT ) | ( 1 << TWSTO ) | ( 1 << TWEN );
}


void TWIWrite( uint8_t u8data )
{
    TWDR = u8data;
    TWCR = ( 1 << TWINT ) | ( 1 << TWEN );

    while ( ( TWCR & ( 1 << TWINT ) )  == 0 );
}


uint8_t TWIReadACK()
{
    TWCR = ( 1 << TWINT ) | ( 1 << TWEN ) | ( 1 << TWEA );

    while ( ( TWCR & ( 1 << TWINT ) ) == 0 );
    
    return TWDR;
}


uint8_t TWIReadNACK()
{
    TWCR = ( 1 << TWINT ) | ( 1 << TWEN );

    while ( ( TWCR & ( 1 << TWINT ) ) == 0 );

    return TWDR;
}


uint8_t TWIGetStatus()
{
    uint8_t status;
    status = TWSR & 0xF8;
    
    return status;
}

ssd1306.c

static int     BUFFER_SIZE   = 128 * 64 / 8;
static uint8_t SCREEN_WIDTH  = 128;

uint8_t screen_buffer[ 128 * 64 / 8 ];


void ssd1306_init( void )
{
    TWIStart();

    TWIWrite( 0x3c << 1 );     // addr?
    TWIWrite( 0xAE );          // display off
    TWIWrite( 0xD9 );          // set display clock division
    TWIWrite( 0x80 );          // the suggested ratio 0x80
    TWIWrite( 0xA8 );          // set multiplex
    TWIWrite( 63   );          // set height
    TWIWrite( 0xD3 );          // set display offset
    TWIWrite( 0x0  );          // no offset
    TWIWrite( 64   );          // line #0 setstartline
    TWIWrite( 0x8D );          // chargepump
    TWIWrite( 0x14 );
    TWIWrite( 0x20 );          // memory d
    TWIWrite( 0x00 );          // 0x0 act like ks0108
    TWIWrite( 161  );          // segremap
    TWIWrite( 0xC8 );          // comscandec
    TWIWrite( 0xDA );          // set com pins
    TWIWrite( 0x12 );
    TWIWrite( 0x81 );          // set contrast
    TWIWrite( 0xFF );
    TWIWrite( 0xD9 );          // 0xd9 set pre-charge
    TWIWrite( 0xF1 );
    TWIWrite( 0xA7 );
    TWIWrite( 0xDB );          
    TWIWrite( 0x40 );
    TWIWrite( 0xA4 );          // display all on resume
    TWIWrite( 0xA7 );          // normal display
    TWIWrite( 0x2E );          // deactivate scroll
    TWIWrite( 0xAF );          // turn on oled panel

    TWIStop();
}


void ssd1306_on()
{
    TWIStart();

    TWIWrite( 0x3c << 1 );
    TWIWrite( 0xA4 );          // display all on resume
    TWIWrite( 0xA6 );          // normal display
    TWIWrite( 0x2E );          // scroll
    TWIWrite( 0xAF );          // turn on oled panel

    TWIStop();
}


void ssd1306_off()
{
    TWIStart();

    TWIWrite( 0x3c << 1 );
    TWIWrite( 0xA4 );          // display all on resume
    TWIWrite( 0xA6 );          // normal display
    TWIWrite( 0x2E );          // deactivate scroll
    TWIWrite( 0xAE );          // display off

    TWIStop();
}


void ssd1306_brightness( int brightness )
{
    TWIStart();

    TWIWrite( 0x3c << 1 );
    TWIWrite( 0x20 );          // memory mode
    TWIWrite( 0x81 );          // contrast
    TWIWrite( brightness );

    TWIStop();
}


void ssd1306_render()
{
    TWIStart();
    TWIWrite( 0x3c << 1 );
    TWIWrite( 0x40 );

    for( int i=0; i<BUFFER_SIZE; i++ )
        TWIWrite( screen_buffer[ i ] );

    TWIStop();
}


void ssd1306_clear()
{
    memset( screen_buffer, 0x00, BUFFER_SIZE );
}


void ssd1306_invert()
{
    TWIStart();

    TWIWrite( 0x3c << 1 );
    TWIWrite( 0xA7 );

    TWIStop();
}


void ssd1306_drawPixel( int x, int y )
{
   screen_buffer[ x + ( y / 8 ) * SCREEN_WIDTH ] |= _BV( ( y % 8 ) ); 
}


void ssd1306_clearPixel( int x, int y )
{
   screen_buffer[ x + ( y / 8 ) * SCREEN_WIDTH ] &= ~_BV( ( y % 8 ) ); 
}

ssd1306_test.c

#define F_CPU 8000000L

#include <avr/io.h>
#include <util/delay.h>
#include <math.h>

#include "protocols/twi.c"
#include "drivers/ssd1306.c"

int main( void )
{
    ssd1306_init();
    ssd1306_on();
 
    int x = 1;

    while( 1 )
    {
        int y1 = sin( x / 2.0 ) * 5;
        int y2 = cos( x / 4.0 ) * 5;

        ssd1306_drawPixel( x, y1 + 15 );
        ssd1306_drawPixel( x, y2 + 40 );
        ssd1306_render();

        x += 1;

        if( x > 127 )
        {
            x = 0;
            ssd1306_clear();
        }

        _delay_ms( 10 );        
    }

}

To build it I use the following commands:

/usr/bin/avr-gcc -g -Os -Wall -mcall-prologues -mmcu=atmega328p   -c -o ssd1306_test.o ssd1306_test.c
/usr/bin/avr-gcc -g -Os -Wall -mcall-prologues -mmcu=atmega328p ssd1306_test.o -o ssd1306_test.obj
/usr/bin/avr-objcopy -R .eeprom -O ihex ssd1306_test.obj ssd1306_test.hex
/usr/bin/avrdude -c linuxgpio -p m328p -e -B 1.0 -U lfuse:w:0xe2:m -U hfuse:w:0xd9:m