Function main

experimental/legacy/transformer/run.cpp:174–247

}

int main() {
  printf("\033[2J\033[1;1H");
  Context ctx = createContext();
  static constexpr size_t seqLen = 24;
  static constexpr size_t batchSize = 1;
  static constexpr size_t modelDim = 4; // 3072;
  static constexpr size_t hiddenWidth = modelDim * 2;
  static constexpr size_t qkvDim = 3; // 256;
  static constexpr size_t nHeads = 8;
  std::mt19937 gen(314);

  Transformer transformer;
  Activations activations;
  KVCache kvcache;
  LOG(kDefLog, kInfo, "Initializing transformer, allocating GPU buffers ...\n");
  createTransformer(ctx, modelDim, qkvDim, nHeads, batchSize, seqLen,
                    hiddenWidth, transformer, activations, kvcache);

  std::array<float, modelDim> inputArr;
  randint(inputArr, gen, -2, 2);
  LOG(kDefLog, kInfo, "%s",
      show<float>(inputArr.data(), 1, modelDim, "Input").c_str());
  Tensor input = createTensor(ctx, Shape{modelDim}, kf32, inputArr.data());

  /* QKV Projection */

  LOG(kDefLog, kInfo, "QKV Projection");
  {
    KernelCode matmul = createMatmul(kShaderMatmul1, /*M*/ batchSize,
                                     /*K*/ modelDim, /*N*/ 3 * qkvDim);
    Kernel qkv = createKernel(
        ctx, matmul, Bindings{input, transformer.qkv, activations.qkv},
        /*nthreads*/ {modelDim, 1, 1});
    std::promise<void> promise;
    std::future<void> future = promise.get_future();
    dispatchKernel(ctx, qkv, promise);
    wait(ctx, future);
    std::array<float, 3 * qkvDim> outputArr;
    toCPU(ctx, activations.qkv, outputArr.data(), sizeof(outputArr));
    LOG(kDefLog, kInfo, "Output: %s",
        show<float>(outputArr.data(), 1, 3 * qkvDim, "QKV Output").c_str());
    std::array<float, 3 * qkvDim> outputRefArr;
    std::array<float, modelDim * 3 * qkvDim> weightsArr;
    toCPU(ctx, transformer.qkv, weightsArr.data(), sizeof(weightsArr));
    ref::matmul_forward_cpu(
        outputRefArr.data(), inputArr.data(), weightsArr.data(), nullptr,
        /* batch */ 1, /* T */ 1, /* C */ modelDim, /* OC */ 3 * qkvDim);
    LOG(kDefLog, kInfo, "Reference Output: %s",
        show<float>(outputRefArr.data(), 1, 3 * qkvDim,
                    "QKV Output (Reference)")
            .c_str());
    LOG(kDefLog, kInfo,
        isclose(outputArr.data(), outputRefArr.data(), 3 * qkvDim) ? "PASS"
                                                                   : "FAIL");
  }

  /* QK Dot Products */