Function testLayerNorm

experimental/legacy/transformer/test_kernels.cpp:168–216

}

void testLayerNorm(Context &ctx) {
  struct LNParam {
    uint32_t N; // check
    uint32_t C;
  };
  constexpr size_t N = 6;
  constexpr size_t C = 3072;
  std::mt19937 gen(31415);
  std::array<float, N * C> inputArr;
  randint(inputArr, gen, 0, 3);
  // range(inputArr);
  std::array<float, N * C> outputArr;
  std::array<float, C> weightArr;
  std::array<float, C> biasArr;
  Tensor input = createTensor(ctx, {N, C}, kf32, inputArr.data());
  LNParam params = {N, C};
  randint(weightArr, gen, 0, 5); // populate randomly
  randint(biasArr, gen, 0, 5);
  Tensor weight = createTensor(ctx, {C}, kf32, weightArr.data());
  Tensor bias = createTensor(ctx, {C}, kf32, biasArr.data());
  Tensor output = createTensor(ctx, {N, C}, kf32, outputArr.data());
  std::promise<void> promise;
  std::future<void> future = promise.get_future();
  Kernel op = createKernel(ctx, {kShaderLayerNorm1, 256, kf32},
                           Bindings{input, weight, bias, output},
                           /* n threads */ {N, 1, 1}, params);
  dispatchKernel(ctx, op, promise);
  wait(ctx, future);
  toCPU(ctx, output, outputArr.data(), sizeof(outputArr));
  LOG(kDefLog, kInfo, "%s",
      show<float, N, C>(inputArr, "LayerNorm Input").c_str());
  LOG(kDefLog, kInfo, "%s",
      show<float, 1, C>(weightArr, "LayerNorm Weight").c_str());
  LOG(kDefLog, kInfo, "%s",
      show<float, 1, C>(biasArr, "LayerNorm Bias").c_str());
  LOG(kDefLog, kInfo, "%s",
      show<float, N, C>(outputArr, "LayerNorm Output").c_str());
  std::array<float, N * C> refOutputArr;
  ref::layernorm_forward_cpu(refOutputArr.data(), inputArr.data(),
                             weightArr.data(), biasArr.data(), N, 1, C);
  LOG(kDefLog, kInfo, "%s",
      show<float, N, C>(refOutputArr,
                        "LayerNorm Reference Implementation Output")
          .c_str());
  bool passed = isclose(outputArr.data(), refOutputArr.data(), N * C);
  assert(passed);
  LOG(kDefLog, kInfo, "LayerNorm passed? %d", passed);
  LOG(kDefLog, kInfo, "Done with LayerNorm Test");
}

void testSoftmax(Context &ctx) {