Function testSoftmax

experimental/legacy/transformer/test_kernels.cpp:218–256

}

void testSoftmax(Context &ctx) {

  struct SoftmaxParam {
    uint32_t N;
    uint32_t C;
  };
  static constexpr size_t B = 6;    // batch size
  static constexpr size_t T = 8;    // token index
  static constexpr size_t C = 3072; // input channels
  std::array<float, B * T * C> inputArr;
  std::array<float, B * T * C> outputArr;
  std::mt19937 gen(31415);
  randint(inputArr, gen, 0, 3);
  Tensor input = createTensor(ctx, {B * T, C}, kf32, inputArr.data());
  Tensor output = createTensor(ctx, {B * T, C}, kf32, outputArr.data());
  LOG(kDefLog, kInfo, "num threads: %d", B * T);
  std::promise<void> promise;
  std::future<void> future = promise.get_future();
  Kernel op = createKernel(
      ctx, {kShaderSoftmax1, 256, kf32}, Bindings{input, output},
      Shape{cdiv(B * T, 256), 1, 1}, SoftmaxParam{B * T, C});
  dispatchKernel(ctx, op, promise);
  wait(ctx, future);
  toCPU(ctx, output, outputArr.data(), sizeof(outputArr));
  LOG(kDefLog, kInfo, "%s",
      show<float, B * T, C>(inputArr, "Softmax Input").c_str());
  LOG(kDefLog, kInfo, "%s",
      show<float, B * T, C>(outputArr, "Softmax Output").c_str());
  std::array<float, B * T * C> refOutputArr;
  ref::softmax_forward_cpu(refOutputArr.data(), inputArr.data(), B * T, C);
  LOG(kDefLog, kInfo, "%s",
      show<float, B * T, C>(refOutputArr, "Softmax reference Output").c_str());

  LOG(kDefLog, kInfo, "number of elements: %d", B * T * C);
  bool passed = isclose(outputArr.data(), refOutputArr.data(), B * T * C);
  assert(passed);
  LOG(kDefLog, kInfo, "Softmax passed? %d", passed);
  LOG(kDefLog, kInfo, "Done with Softmax Test");
}

void testAttention(Context &ctx) {
  static constexpr size_t B = 6;