text: /*
   * Copyright (c) Facebook, Inc. and its affiliates.
   * All rights reserved.
   * This source code is licensed under the BSD-style license found in the
   * LICENSE file in the root directory of this source tree.
   */
  #include <assert.h>
  #include <cpuid.h>
  #include <stdlib.h>
  #include <string.h>
  #include <algorithm>
  #include <array>
  #include <fstream>
  #include <functional>
  #include <iostream>
  #include <map>
  #include <string>
  
  using namespace std;
  
  void addi(ofstream& of, string i, bool disable = false) {
    if (disable == false)
      of << "      \"" + i + "\\t\\n\"" + "\n";
  }
  
  struct ISA {
    unsigned avx; // 1, 2 or 3
    string name;
    vector<vector<unsigned>> shapes;
  };
  
  int main() {
    bool iaca = false;
    bool disable = false;
  
    bool fixedA = true, fixedB = true, fixedC = true;
  
    int eax, ebx, ecx, edx;
    __cpuid(1 /* ecx = vendor string */, eax, ebx, ecx, edx);
    printf("FC16 is %s supported\n", ((ecx & bit_F16C) ? " " : "not"));
  
    string comma = ",";
  
    vector<ISA> isa = {
        // {1, "AVX", {{4, 1, 0}, {4, 2, 0}, {4, 3, 0}, {3, 1, 0}, {3, 2, 0}, {3,
        // 3, 0}}},
        {2,
         "AVX2",
         {
             // 4x3 register layout
             // {1, 3, 0},
             // {2, 3, 0},
             // {3, 3, 0},
             // {4, 3, 0},
  
             // 6x2 register layout
             {1, 2, 0},
             {2, 2, 0},
             {3, 2, 0},
             {4, 2, 0},
             {5, 2, 0},
             {6, 2, 0},
  
             // 14x1 register layout
             // {1, 1, 0},
             // {2, 1, 0},
             // {3, 1, 0},
             // {4, 1, 0},
             // {5, 1, 0},
             // {6, 1, 0},
             // {7, 1, 0},
             // {8, 1, 0},
             // {9, 1, 0},
             // {10, 1, 0},
             // {11, 1, 0},
             // {12, 1, 0},
             // {13, 1, 0},
             // {14, 1, 0},
         }},
        {3,
         "AVX512",
         {
             // 14x2 register layout
             {1, 2, 0},
             {2, 2, 0},
             {3, 2, 0},
             {4, 2, 0},
             {5, 2, 0},
             {6, 2, 0},
             {7, 2, 0},
             {8, 2, 0},
             {9, 2, 0},
             {10, 2, 0},
             {11, 2, 0},
             {12, 2, 0},
             {13, 2, 0},
             {14, 2, 0},
         }}};
  
    for (auto s : isa) {
      string isa_file_name = s.avx <= 2 ? "Avx2" : "Avx512";
  
      // open all files
      ofstream srcfile;
      srcfile.open("FbgemmFP16UKernels" + isa_file_name + ".cc");
      srcfile
          << "/*\n"
             " * Copyright (c) Facebook, Inc. and its affiliates.\n"
             " * All rights reserved.\n"
             " * This source code is licensed under the BSD-style license found in the\n"
             " * LICENSE file in the root directory of this source tree.\n"
             " */\n";
      srcfile << "#include \"FbgemmFP16UKernels" + isa_file_name + ".h\"\n\n";
      srcfile << "namespace fbgemm {\n\n";
      if (iaca) {
        srcfile << "#include \"iacaMarks.h\"\n";
      }
  
      ofstream hdrfile;
      hdrfile.open("FbgemmFP16UKernels" + isa_file_name + ".h");
      hdrfile
          << "/*\n"
             " * Copyright (c) Facebook, Inc. and its affiliates.\n"
             " * All rights reserved.\n"
             " * This source code is licensed under the BSD-style license found in the\n"
             " * LICENSE file in the root directory of this source tree.\n"
             " */\n";
  
      hdrfile << "#pragma once\n";
      hdrfile << "#include <cstdint>\n";
      if (s.avx == 3) {
        hdrfile << "#include \"FbgemmFP16UKernelsAvx2.h\"\n";
      }
      hdrfile << "#include \"fbgemm/Types.h\"\n\n";
      hdrfile << "namespace fbgemm {\n\n";
      if (s.avx == 2) {
        hdrfile << "using fp16 = float16;\n";
        hdrfile << "using fp32 = float;\n";
        hdrfile
            << "struct GemmParams {\n  uint64_t k;\n  float* A;\n  const fp16* B;\n"
               "  float* beta;\n  uint64_t accum;\n  float* C;\n  uint64_t ldc;\n"
               "  uint64_t b_block_cols;\n  uint64_t b_block_size;\n};\n";
      }
  
      map<string, string> fptr_typedef;
      fptr_typedef["fp16"] = "";
      fptr_typedef["fp32"] = "";
  
      unsigned labelId = 0;
  
      bool fixedA = false, fixedB = false, fixedC = false;
      bool fp16 = true;
  
      vector<vector<unsigned>>& ukernel_shape = s.shapes;
  
      vector<string> funcname(ukernel_shape.size()),
          fheader(ukernel_shape.size());
      string fargs;
  
      string B_type = ((fp16) ? "fp16" : "fp32");
      string prefix = s.name + /*"_" + B_type */ +"_" + "fA" + to_string(fixedA) +
          "fB" + to_string(fixedB) + "fC" + to_string(fixedC);
      cout << "Generating code for " << s.name << " " << B_type << "\n";
  
      string vec_reg_prefix = s.avx <= 2 ? "ymm" : "zmm";
      int num_vec_regs = s.avx <= 2 ? 16 : 32;
      int vec_len_in_bytes = s.avx <= 2 ? 32 : 64;
  
      for (unsigned k = 0; k < ukernel_shape.size(); k++) {
        printf("shape: %d x %d * 32\n", ukernel_shape[k][0], ukernel_shape[k][1]);
  
        string p1 = "GemmParams* gp";
  
        funcname[k] = "gemmkernel_" + to_string(ukernel_shape[k][0]) + "x" +
            to_string(ukernel_shape[k][1]) + "_";
        funcname[k] += prefix;
  
        fargs = "(" + p1 + ")";
  
        fheader[k] = "void __attribute__((noinline)) " + funcname[k] + fargs;
        srcfile << fheader[k] << " {\n";
  
        unsigned last_free_vecreg = 0;
        // produce register block of C
        vector<vector<string>> vCtile(ukernel_shape[k][0]);
        for (auto r = 0; r < ukernel_shape[k][0]; r++)
          for (auto c = 0; c < ukernel_shape[k][1]; c++) {
            vCtile[r].push_back(vec_reg_prefix + to_string(last_free_vecreg));
            last_free_vecreg++;
          }
        assert(last_free_vecreg <= num_vec_regs - 2);
  
        string vAtmp = vec_reg_prefix + to_string(last_free_vecreg++);
        // produce register block of B col
        vector<string> vBcol(ukernel_shape[k][1]);
  
        for (auto c = 0; c < ukernel_shape[k][1]; c++) {
          vBcol[c] = (vec_reg_prefix + to_string(last_free_vecreg));
          last_free_vecreg++;
        }
  
        assert(last_free_vecreg <= num_vec_regs);
  
        srcfile << "  asm volatile(\n";
  
        srcfile << "#if !defined(__clang__)"
                << "\n";
        addi(srcfile, "mov r14, %[gp]");
        srcfile << "#else\n";
        addi(srcfile, "mov %[gp], %%r14");
        addi(srcfile, ".intel_syntax noprefix");
        srcfile << "#endif\n";
  
        srcfile << "\n      // Copy parameters\n";
        srcfile << "      // k\n";
        addi(srcfile, "mov r8, [r14 + 0]");
        srcfile << "      // A\n";
        addi(srcfile, "mov r9, [r14 + 8]");
        srcfile << "      // B\n";
        addi(srcfile, "mov r10, [r14 + 16]");
        srcfile << "      // beta\n";
        addi(srcfile, "mov r15, [r14 + 24]");
        srcfile << "      // accum\n";
        addi(srcfile, "mov rdx, [r14 + 32]");
        srcfile << "      // C\n";
        addi(srcfile, "mov r12, [r14 + 40]");
        srcfile << "      // ldc\n";
        addi(srcfile, "mov r13, [r14 + 48]");
        srcfile << "      // b_block_cols\n";
        addi(srcfile, "mov rdi, [r14 + 56]");
        srcfile << "      // b_block_size\n";
        addi(srcfile, "mov rsi, [r14 + 64]");
        srcfile << "      // Make copies of A and C\n";
        addi(srcfile, "mov rax, r9");
        addi(srcfile, "mov rcx, r12");
        srcfile << "\n";
  
        addi(srcfile, "mov rbx, 0");
  
        string exitlabel = "L_exit%=";
        string label2 = "loop_outter%=";
        addi(srcfile, label2 + ":");
        addi(srcfile, "mov r14, 0");
  
        // set all vCtile regs to zeros
        for (auto r = 0; r < vCtile.size(); r++) {
          for (auto c = 0; c < vCtile[r].size(); c++) {
            addi(
                srcfile,
                "vxorps " + vCtile[r][c] + "," + vCtile[r][c] + "," +
                    vCtile[r][c]);
          }
        }
  
        // start marker
        if (iaca) {
          addi(srcfile, "mov ebx, 111");
          addi(srcfile, ".byte 0x64, 0x67, 0x90");
        }
  
        srcfile << "\n";
  
        srcfile << "\n";
        string label = "loop_inner%=";
        addi(srcfile, label + ":");
        srcfile << "\n";
  
        for (int c = 0; c < vCtile[0].size(); c++) {
          addi(
              srcfile,
              "vcvtph2ps " + vBcol[c] + "," + (s.avx <= 2 ? "XMM" : "YMM") +
                  "WORD PTR [r10 + " + to_string(vec_len_in_bytes / 2 * c) + "]");
        }
  
        for (int r = 0; r < vCtile.size(); r++) {
          addi(
              srcfile,
              "vbroadcastss " + vAtmp + ",DWORD PTR [r9+" + to_string(4 * r) +
                  "]");
          for (int c = 0; c < vCtile[0].size(); c++) {
            addi(
                srcfile,
                "vfmadd231ps " + vCtile[r][c] + "," + vBcol[c] + "," + vAtmp);
          }
        }
  
        addi(
            srcfile,
            "add r9," + to_string(4 * ukernel_shape[k][0]),
            fixedA); // move A ptr
  
        addi(
            srcfile,
            "add r10," + to_string(vec_len_in_bytes / 2 * ukernel_shape[k][1]),
            fixedA); // move A ptr
  
        addi(srcfile, "inc r14");
        addi(srcfile, "cmp r14, r8");
        addi(srcfile, "jl " + label);
  
        srcfile << "\n";
  
        addi(srcfile, exitlabel + ":");
  
        // addi(srcfile, "add r10, rsi");
        srcfile << "\n";
  
        // end marker
        if (iaca) {
          addi(srcfile, "mov ebx, 222");
          addi(srcfile, ".byte 0x64, 0x67, 0x90");
        }
  
        addi(srcfile, "cmp rdx, 1");
        addi(srcfile, "je L_accum%=");
        srcfile << "      // Dump C\n";
  
        for (auto r = 0; r < vCtile.size(); r++) {
          for (auto c = 0; c < vCtile[r].size(); c++) {
            addi(
                srcfile,
                "vmovups " + vec_reg_prefix + "word PTR [r12 + " +
                    to_string(vec_len_in_bytes * c) + "], " + vCtile[r][c],
                fixedC);
          }
          addi(srcfile, "add r12, r13", fixedC); // move C ptr
        }
        addi(srcfile, "jmp L_done%=");
  
        srcfile << "\n";
        addi(srcfile, "L_accum%=:");
        srcfile << "      // Dump C with accumulate\n";
  
        string r_spare =
            vec_reg_prefix + to_string(num_vec_regs - (s.avx == 1 ? 2 : 1));
        string r_last = vec_reg_prefix + to_string(num_vec_regs - 1);
        addi(
            srcfile,
            "vbroadcastss " + r_spare + string(",DWORD PTR [r15]"),
            fixedC);
        // store out C
        for (auto r = 0; r < vCtile.size(); r++) {
          for (auto c = 0; c < vCtile[r].size(); c++) {
            switch (s.avx) {
              case 1:
                addi(
                    srcfile,
                    string("vmulps " + r_last + ", ") + r_spare + comma +
                        "YMMWORD PTR [r12 + " + to_string(vec_len_in_bytes * c) +
                        "]",
                    fixedC);
                addi(
                    srcfile,
                    "vaddps " + vCtile[r][c] + "," + vCtile[r][c] + "," + r_last,
                    fixedC);
                break;
              case 2:
              case 3:
                addi(
                    srcfile,
                    "vfmadd231ps " + vCtile[r][c] + "," + r_spare + "," +
                        vec_reg_prefix + "word PTR [r12 + " +
                        to_string(vec_len_in_bytes * c) + "]",
                    fixedC);
                break;
              default:
                assert(0);
            }
            addi(
                srcfile,
                "vmovups " + vec_reg_prefix + "word PTR [r12 + " +
                    to_string(vec_len_in_bytes * c) + "], " + vCtile[r][c],
                fixedC);
          }
          addi(srcfile, "add r12, r13", fixedC); // move C ptr
        }
  
        srcfile << "\n";
        addi(srcfile, "L_done%=:");
  
        srcfile << "\n      // next outer iteration\n";
        // C
        addi(
            srcfile,
            "add rcx, " + to_string(vec_len_in_bytes * ukernel_shape[k][1]),
            fixedC);
        addi(srcfile, "mov r12, rcx", fixedC);
        // A
        addi(srcfile, "mov r9, rax");
  
        addi(srcfile, "inc rbx");
        addi(srcfile, "cmp rbx, rdi");
        addi(srcfile, "jl " + label2);
  
        // output
        srcfile << "      :\n";
        // input
        srcfile << "      : [gp] \"rm\"(gp)\n";
  
        // clobbered
        srcfile << "      : \"r8\",\n        \"r9\",\n        \"r10\",\n"
                   "        \"r11\",\n        \"r15\",\n        \"r13\",\n"
                   "        \"r14\",\n        \"rax\",\n        \"rcx\",\n"
                   "        \"rdx\",\n        \"rsi\",\n        \"rdi\",\n"
                   "        \"rbx\",\n        \"r12\",\n"
                   "        \"memory\");\n";
        srcfile << "}\n";
      }
  
      for (unsigned k = 0; k < ukernel_shape.size(); k++) {
        hdrfile << fheader[k] << ";\n";
      }
  
      fptr_typedef[B_type] = "typedef void (*funcptr_" + B_type + ")" + fargs;
  
      srcfile << "\n} // namespace fbgemm\n";
      srcfile.close();
  
      hdrfile << fptr_typedef["fp16"] << ";\n";
      hdrfile << fptr_typedef["fp32"] << ";\n";
      hdrfile << "\n} // namespace fbgemm\n\n";
      hdrfile.close();
    } // isa
  }
  