RISC-VにおけるCrypto命令 (Crypto Extension)

やっとこさVector Extensionも1.0になり、そろそろいろんなExtensionを実装したSoCも出てくるとは思うけど、年末の移動時間のなかでいろいろ現在のCrypto Extensionを勉強してた

Crypto ExtensioもScalar命令版とVector命令版の２つがあって、Scalar版は組み込みとか用途向けの、レジスタとかを増やしたくない実装用で、Vectorはその名の通りVectorレジスタを使ったもの。試しにAESのEncryptionを実装してみる

Scalar版

Scalar版は以下のようになる。

void
riscv64zkn_aes_encrypt_block_128(uint64_t* expandedKey,
                                 uint8_t *output,
                                 const uint8_t *input)
{
  uint64_t state0 = *((uint64_t *)input);
  uint64_t state1 = *((uint64_t *)(input + 8));

  // Add round key 0 to initial state
  state0 = state0 ^ *expandedKey++;
  state1 = state1 ^ *expandedKey++;

  for (int i = 1; i < 10; r++) {
    uint64_t c0 = __riscv_aes64esm(state0, state1);
    uint64_t c1 = __riscv_aes64esm(state1, state0);
    state0 = c0 ^ *expandedKey++;
    state1 = c1 ^ *expandedKey++;
  }

  // Final round
  uint64_t c0 = __riscv_aes64es(state0, state1);
  uint64_t c1 = __riscv_aes64es(state1, state0);
  state0 = c0 ^ *expandedKey++;
  state1 = c1 ^ *expandedKey;

  *((uint64_t *)output) = state0;
  *((uint64_t *)(output + 8)) = state1;
}

Vector版

Vector版。まだVector Crypto ExtensionのIntrinics命令はStableじゃないので、インラインアセンブラ使ってる。

static vuint32m4_t
vaesz_vs(vuint32m4_t rd, vuint32m4_t vs2)
{
  __asm__("vaesz.vs %0, %1" : "+vr"(rd) : "vr"(vs2));
  return rd;
}

static vuint32m4_t
vaesem_vs(vuint32m4_t rd, vuint32m4_t vs2)
{
  __asm__("vaesem.vs %0, %1" : "+vr"(rd) : "vr"(vs2));
  return rd;
}

static vuint32m4_t
vaesef_vs(vuint32m4_t rd, vuint32m4_t vs2)
{
  __asm__("vaesef.vs %0, %1" : "+vr"(rd) : "vr"(vs2));
  return rd;
}

SECStatus
riscv64zvkn_aes_encrypt_block_128(uint32_t* expandedKey,
                                  uint8_t *output,
                                  const uint8_t *input)
{
  size_t vl = __riscv_vsetvl_e32m4(4);
  vuint32m4_t state = __riscv_vle32_v_u32m4((const uint32_t *)input, vl);

  // Add round key 0 to initial state
  vuint32m4_t K = __riscv_vle32_v_u32m4(expandedKey, vl);
  expandedKey += 4;
  state = vaesz_vs(state, K);

  for (int i = 1; i < 10; r++) {
    K = __riscv_vle32_v_u32m4(expandedKey, vl);
    expandedKey += 4;
    state = vaesem_vs(state, K);
  }

  // Final round
  K = __riscv_vle32_v_u32m4(expandedKey, vl);
  state = vaesef_vs(state, K);
  __riscv_vse32_v_u32m4((uint32_t *)output, state, vl);
}

Conclusion

• RISC-Vの場合Load/Restoreをまとめないとvlenの長さを指定するvsetvliが大量に利用することになるので、速さに直結するかもということと、コンパイラ (gcc、LLVM) の最適化もまだまだっぽい
• 両方ともx86のAES-NIに似た感じなので、AES-NIの経験値があれば、難しくないかも