Works all 24 rounds without paddin (essentially just keccakf)

gpu_keccak.js

"use strict";

const KECCAK_ROUNDS = 24;

var input = new Uint32Array([
    1063097118, 728099034, 2853146927, 1330851092, 209328888,
    904790539, 786070775, 1411619771, 1176535054, 1708240850,
    1280969794, 3252320020, 3814437729, 900990975, 3404197392,
    3203136532, 317488412, 1753212407, 1201418786, 3729619196,
    1635356607, 2938933362, 3579979232, 3308551012, 3620133985,
    1264321397, 3897539483, 2395261390, 472238529, 742831769,
    3333799317, 3040321574, 2855965447, 170043452, 2478026733,
    3378843531, 3079917733, 1032759894, 3218255483, 4024074026,
    3535953065, 3283636176, 2658873612, 31102196, 3587910756,
    1426793143, 2221768042, 409214893, 24442995, 146832452
]);

var theta_chi = new Int8Array([ //RGBA8I
  4, 1,  1,   2,
  0, 2,  1,   2,
  1, 3,  1,   2,
  2, 4,  1,  -3,
  3, 0, -4,  -3
]);

var rho_pi = new Int8Array([ // RGBA8I
  // red is the index to rotate
  // green is the number of bits to rotate
  // blue is red % 5 [needed for indexing into theta_chi]
  // alpha is floor(green / 16) [needed for calculating 16 bit rotation]
  0, 0, 0, 0,
  6, 44, 1, 2,
  12, 43, 2, 2,
  18, 21, 3, 1,
  24, 14, 4, 0,
  3, 28, 3, 1,
  9, 20, 4, 1,
  10, 3, 0, 0,
  16, 45, 1, 2,
  22, 61, 2, 3,
  1, 1, 1, 0,
  7, 6, 2, 0,
  13, 25, 3, 1,
  19, 8, 4, 0,
  20, 18, 0, 1,
  4, 27, 4, 1,
  5, 36, 0, 2,
  11, 10, 1, 0,
  17, 15, 2, 0,
  23, 56, 3, 3,
  2, 62, 2, 3,
  8, 55, 3, 3,
  14, 39, 4, 2,
  15, 41, 0, 2,
  21, 2, 1, 0
]);

var keccakf_rndc  = new Uint16Array([
  1, 0, 0, 0,
  32898, 0, 0, 0,
  32906, 0, 0, 32768,
  32768, 32768, 0, 32768,
  32907, 0, 0, 0,
  1, 32768, 0, 0,
  32897, 32768, 0, 32768,
  32777, 0, 0, 32768,
  138, 0, 0, 0,
  136, 0, 0, 0,
  32777, 32768, 0, 0,
  10, 32768, 0, 0,
  32907, 32768, 0, 0,
  139, 0, 0, 32768,
  32905, 0, 0, 32768,
  32771, 0, 0, 32768,
  32770, 0, 0, 32768,
  128, 0, 0, 32768,
  32778, 0, 0, 0,
  10, 32768, 0, 32768,
  32897, 32768, 0, 32768,
  32896, 0, 0, 32768,
  1, 32768, 0, 0,
  32776, 32768, 0, 32768,
]);

var vertexShaderSource = `#version 300 es

in vec2 a_position;
in vec2 a_textureCoordinates;

uniform vec2 u_resolution;

out vec2 v_textureCoordinates;

// all shaders have a main function
void main() {

  // convert the position from pixels to 0.0 to 1.0
  vec2 zeroToOne = a_position / u_resolution;

  // convert from 0->1 to 0->2
  vec2 zeroToTwo = zeroToOne * 2.0;

  // convert from 0->2 to -1->+1 (clipspace)
  vec2 clipSpace = zeroToTwo - 1.0;

  // must set gl_Position in every vertex shader
  gl_Position = vec4(clipSpace * vec2(1, -1), 0, 1);

  v_textureCoordinates = a_textureCoordinates;
}
`;

var fragmentShaderSource = `#version 300 es

// fragment shaders don't have a default precision so we need
// to pick one. mediump is a good default. It means "medium precision"
precision highp float;
precision highp int;
precision highp usampler2D;
precision highp sampler2D;
precision highp isampler2D; 

uniform usampler2D input_data;
uniform isampler2D theta_chi;
uniform isampler2D rho_pi;
uniform usampler2D keccakf_rndc;

uniform int round;

in vec2 v_textureCoordinates;
out uvec4 outColor;

// for rotations below 16 bits
uvec4 rotl_simple(uvec4 num, int shift) {
  uint hihi = (num[3] << shift) | (num[2] >> (16-shift));
  uint hilo = (num[2] << shift) | (num[1] >> (16-shift));
  uint lohi = (num[1] << shift) | (num[0] >> (16-shift));
  uint lolo = (num[0] << shift) | (num[3] >> (16-shift));
  uvec4 result = uvec4(lolo, lohi, hilo, hihi);

  // because shader is not aware of uint16 representation we & with 0xffff
  result &= 65535u;

  return result;
}

uvec4 rotl(uvec4 num, int shift, int movs) {
  for(int i=0; i<movs; i++) {
    num = uvec4(num[3], num[0], num[1], num[2]);
  }
  shift -= (movs*16);
  return rotl_simple(num, shift);
}

uvec4 trp(int i) {
  ivec4 rp_idxs = texelFetch(rho_pi, ivec2(i, 0), 0).rgba;
  ivec2 theta_idxs = texelFetch(theta_chi, ivec2(rp_idxs.b, 0), 0).rg;

  uvec4 xor = texelFetch(input_data, ivec2(theta_idxs.r, 0), 0) ^
              texelFetch(input_data, ivec2(theta_idxs.r+5, 0), 0) ^
              texelFetch(input_data, ivec2(theta_idxs.r+10, 0), 0) ^
              texelFetch(input_data, ivec2(theta_idxs.r+15, 0), 0) ^
              texelFetch(input_data, ivec2(theta_idxs.r+20, 0), 0);

  uvec4 t = texelFetch(input_data, ivec2(theta_idxs.g, 0), 0) ^
            texelFetch(input_data, ivec2(theta_idxs.g+5, 0), 0) ^
            texelFetch(input_data, ivec2(theta_idxs.g+10, 0), 0) ^
            texelFetch(input_data, ivec2(theta_idxs.g+15, 0), 0) ^
            texelFetch(input_data, ivec2(theta_idxs.g+20, 0), 0);

  uvec4 pre_rot = texelFetch(input_data, ivec2(rp_idxs.r, 0), 0) ^
                  (xor ^ rotl_simple(t, 1));

  return rotl(pre_rot, rp_idxs.g, rp_idxs.a);
}

void main() {
  // use i.x to store thread index (i.y is always 0)
  // use thread index to lookup j (from rho pi) for the theta portion of lookups
  // use thread index to lookup rho pi and chi portions.

  int i = int(gl_FragCoord.x);
  int s = i%5;

  ivec2 chi_idxs = texelFetch(theta_chi, ivec2(s, 0), 0).ba;
  uvec4 post_rho_pi = trp(i);
  uvec4 prp1 = trp(i+chi_idxs[0]);
  uvec4 prp2 = trp(i+chi_idxs[1]);

  post_rho_pi ^= ((~prp1) & prp2);

  if (i == 0) post_rho_pi ^= texelFetch(keccakf_rndc, ivec2(round, 0), 0);

  outColor = post_rho_pi;
}
`;

function main() {
  
  // https://webgl2fundamentals.org/webgl/lessons/webgl-image-processing-continued.html
  // and use twgl for setting uniforms, attributes, programInfo, etc. Seems good.
  
  var gl = setupWebGLFromCanvas(document.getElementById('c')); // util func
  var ext = gl.getExtension('EXT_color_buffer_float');
  if (!ext) console.log("ERROR EXTENSION");
  var program = createProgramWithShaders(gl, vertexShaderSource, fragmentShaderSource);

  var vertexArrayObject = getVertexArrayObject(gl);
  
  var textureVertices = new Float32Array([
    0.0,  0.0,
    1.0,  0.0,
    0.0,  1.0,
    0.0,  1.0,
    1.0,  0.0,
    1.0,  1.0
  ]);

  var size = 2;          // 2 components per iteration
  var type = gl.FLOAT;   // the data is 32bit floats
  var normalize = false; // don't normalize the data
  var stride = 0;        // 0 = move forward size * sizeof(type) each iteration to get the next position
  var offset = 0;        // start at the beginning of the buffer
  var textureCoordinatesAttributeLocation = gl.getAttribLocation(program, "a_textureCoordinates");
  setupVerticesAndEnableAttribArray(gl, textureVertices, textureCoordinatesAttributeLocation, size, type, normalize, stride, offset);

  var inputTexture = createAndSetupMultiTexture(gl, 0);

  var mipLevel = 0;                 // the largest mip
  var internalFormat = gl.RGBA16UI;   // format we want in the texture
  var srcFormat = gl.RGBA_INTEGER;    // format of data we are supplying
  var srcType = gl.UNSIGNED_SHORT;    // type of data we are supplying
  var width = 25;
  var height = 1;
  var border = 0;
  gl.pixelStorei(gl.UNPACK_ALIGNMENT, 1);
  var sixteen_input = new Uint16Array(input.buffer);
  gl.texImage2D(gl.TEXTURE_2D, mipLevel, internalFormat, width, height, border, srcFormat, srcType, sixteen_input);

  var thetaChiTexture = createAndSetupMultiTexture(gl, 1);
  internalFormat = gl.RGBA8I;
  srcFormat = gl.RGBA_INTEGER;
  srcType = gl.BYTE;
  width = 5;
  gl.texImage2D(gl.TEXTURE_2D, mipLevel, internalFormat, width, height, border, srcFormat, srcType, theta_chi);

  var rhoPiTexture = createAndSetupMultiTexture(gl, 2);
  internalFormat = gl.RGBA8I;
  srcFormat = gl.RGBA_INTEGER;
  srcType = gl.BYTE
  width = 25;
  gl.texImage2D(gl.TEXTURE_2D, mipLevel, internalFormat, width, height, border, srcFormat, srcType, rho_pi);

  // best to make this a uniform which changes every round of array drawing
  var keccakfRndcTexture = createAndSetupMultiTexture(gl, 3);
  internalFormat = gl.RGBA16UI;
  srcFormat = gl.RGBA_INTEGER;
  srcType = gl.UNSIGNED_SHORT;
  width = 24;
  gl.texImage2D(gl.TEXTURE_2D, mipLevel, internalFormat, width, height, border, srcFormat, srcType, keccakf_rndc);

  setupCanvas(gl);

  var outputTexture = createAndSetupMultiTexture(gl, 4);
  internalFormat = gl.RGBA16UI;
  srcFormat = gl.RGBA_INTEGER;
  srcType = gl.UNSIGNED_SHORT;
  width = 25;
  gl.texImage2D(gl.TEXTURE_2D, mipLevel, internalFormat, width, height, border, srcFormat, srcType, null);
  
  // Tell it to use our program (pair of shaders)
  gl.useProgram(program);
  
  // Bind the attribute/buffer set we want.
  gl.bindVertexArray(vertexArrayObject);

  // play with moving this block around
  var inputLocation = gl.getUniformLocation(program, "input_data");
  var thetaChiLocation = gl.getUniformLocation(program, "theta_chi");
  var rhoPiLocation = gl.getUniformLocation(program, "rho_pi");
  var keccakfRndcLocation = gl.getUniformLocation(program, "keccakf_rndc");

  gl.uniform1i(inputLocation, 0);
  gl.uniform1i(thetaChiLocation, 1);
  gl.uniform1i(rhoPiLocation, 2);
  gl.uniform1i(keccakfRndcLocation, 3);

  var resolutionLocation = gl.getUniformLocation(program, "u_resolution");
  gl.uniform2f(resolutionLocation, gl.canvas.width, gl.canvas.height);

  var vertices = createVertexArrayFromImage(width, height);
  var positionAttributeLocation = gl.getAttribLocation(program, "a_position");
  setupVerticesAndEnableAttribArray(gl, vertices, positionAttributeLocation, size, type, normalize, stride, offset);
  
  // draw
  var primitiveType = gl.TRIANGLES;
  var offset = 0;
  var count = 6;

  // framebuffer for ping-ponging
  var framebuffer = gl.createFramebuffer();
  var attachmentPoint = gl.COLOR_ATTACHMENT0;
  
  for(var i=0; i<KECCAK_ROUNDS; i++) { // make a better construct for this

    if (i%2==1){
      bindTextures(gl, [outputTexture, thetaChiTexture, rhoPiTexture, keccakfRndcTexture]);
      bindTextureToFramebuffer(gl, inputTexture, attachmentPoint, framebuffer);
    } else {
      bindTextures(gl, [inputTexture, thetaChiTexture, rhoPiTexture, keccakfRndcTexture]);
      bindTextureToFramebuffer(gl, outputTexture, attachmentPoint, framebuffer);
    }
    var roundLocation = gl.getUniformLocation(program, "round");
    gl.uniform1i(roundLocation, i);
    gl.drawArrays(primitiveType, offset, count);
  }

  readRenderedOutput(gl, width, height);
}

function createProgramWithShaders(gl, vertexShaderSource, fragmentShaderSource) {
  var vertexShader = createShader(gl, gl.VERTEX_SHADER, vertexShaderSource);
  var fragmentShader = createShader(gl, gl.FRAGMENT_SHADER, fragmentShaderSource);
  
  return createProgram(gl, vertexShader, fragmentShader);
}

function setupCanvas(gl) {
  
  // Tell WebGL how to convert from clip space to pixels
  gl.viewport(0, 0, gl.canvas.width, gl.canvas.height);

  // Clear the canvas
  gl.clearColor(0, 0, 0, 0);
  gl.clear(gl.COLOR_BUFFER_BIT);
}

function createVertexArrayFromImage(width, height) {
  var x1 = 0;
  var x2 = width;
  var y1 = 0;
  var y2 = height;
  return new Float32Array([
    x1, y1,
    x2, y1,
    x1, y2,
    x1, y2,
    x2, y1,
    x2, y2
  ]);
}

function readRenderedOutput(gl, width, height) {
  
  // 4 because the fragment shader declares output colors as vec4 (aka RGBA)
  var output_buffer = new ArrayBuffer(gl.drawingBufferWidth * gl.drawingBufferHeight * 4 * 4);
  var output_view = new Uint32Array(output_buffer);

  switch(gl.checkFramebufferStatus(gl.FRAMEBUFFER)) {
    case gl.FRAMEBUFFER_COMPLETE:
      gl.readBuffer(gl.COLOR_ATTACHMENT0); // WARNING: CHANGE THIS!!!!
      gl.readPixels(0, 0, gl.drawingBufferWidth, gl.drawingBufferHeight, gl.RGBA_INTEGER, gl.UNSIGNED_INT, output_view);
      var conversion = new Uint16Array(output_view);
      var output = new Uint32Array(conversion.buffer);
      console.log(output);
      break;
    case gl.FRAMEBUFFER_INCOMPLETE_ATTACHMENT:
      console.log("ERROR: FRAMEBUFFER_INCOMPLETE_ATTACHMENT");
      break;
    case gl.FRAMEBUFFER_INCOMPLETE_DIMENSIONS:
      console.log("ERROR: FRAMEBUFFER_INCOMPLETE_DIMENSIONS");
      break;
    case gl.FRAMEBUFFER_INCOMPLETE_MISSING_ATTACHMENT:
      console.log("ERROR: FRAMEBUFFER_INCOMPLETE_MISSING_ATTACHMENT");
      break;
    case gl.FRAMEBUFFER_UNSUPPORTED:
      console.log("ERROR: FRAMEBUFFER_UNSUPPORTED");
      break;
  }
  
}

main();