commit 9e98bcdbee19371e04396fd6cfbbc3ab7d4f9930
author bjacob <benoitjacob@google.com> Mon Feb 07 23:04:28 2022 -0500
committer GitHub <noreply@github.com> Mon Feb 07 20:04:28 2022 -0800
tree 44082ff396ff2ce6d69f0e5dcd42f9715d1fa275
parent 3944dbb147477cc5a5d65579437a6d3ca1b7442a

Generalize the asm kernel pattern. (#8248)

This makes MMTCustomKernelPattern generic in what kernel it's
generating, in a way that should suffice for the next few kernels that
we need to add (first motivation: float32/aarch64).

MMTCustomKernelPattern now takes the kernel-specific details from a
MMTKernel struct, including the asm code and information about the shape
and data types. The inline asm constraints string becomes
auto-generated.

All this automation had the effect of making it "too magical" for the
inline asm to continue to refer to its register operands by fixed
indices like "$12" as the llvm inline asm syntax allows, so we introduce
a pre-processing step allowing us to write kernels using placeholders
like "$(lhs:0)" for the 0-th register of the LHS, and that gets
"allocated" as the pre-processed code and the constraints string get
co-generated.

This makes for a self-contained way to describe kernels. The current
aarch64 +dotprod i8 kernel becomes:

```
// i8*i8->i32 kernel for Aarch64 NEON +dotprod
MMTKernel MMTKernel_8x4x8_i8i8i32_Aarch64Dotprod_InlineAsm() {
  MMTKernel kernel;
  kernel.arch = CustomKernelTargetArch::Aarch64;
  kernel.m0 = 8;
  kernel.k0 = 4;
  kernel.n0 = 8;
  kernel.lhsType = MMTKernel::ScalarType::I8;
  kernel.rhsType = MMTKernel::ScalarType::I8;
  kernel.accType = MMTKernel::ScalarType::I32;
  kernel.registerBitWidth = PowerOfTwo(128);
  kernel.implAsm = R"ASM(
      sdot $(acc:0).4s, $(rhs:0).16b, $(lhs:0).4b[0]
      sdot $(acc:1).4s, $(rhs:1).16b, $(lhs:0).4b[0]
      sdot $(acc:2).4s, $(rhs:0).16b, $(lhs:0).4b[1]
      sdot $(acc:3).4s, $(rhs:1).16b, $(lhs:0).4b[1]
      sdot $(acc:4).4s, $(rhs:0).16b, $(lhs:0).4b[2]
      sdot $(acc:5).4s, $(rhs:1).16b, $(lhs:0).4b[2]
      sdot $(acc:6).4s, $(rhs:0).16b, $(lhs:0).4b[3]
      sdot $(acc:7).4s, $(rhs:1).16b, $(lhs:0).4b[3]
      sdot $(acc:8).4s, $(rhs:0).16b, $(lhs:1).4b[0]
      sdot $(acc:9).4s, $(rhs:1).16b, $(lhs:1).4b[0]
      sdot $(acc:10).4s, $(rhs:0).16b, $(lhs:1).4b[1]
      sdot $(acc:11).4s, $(rhs:1).16b, $(lhs:1).4b[1]
      sdot $(acc:12).4s, $(rhs:0).16b, $(lhs:1).4b[2]
      sdot $(acc:13).4s, $(rhs:1).16b, $(lhs:1).4b[2]
      sdot $(acc:14).4s, $(rhs:0).16b, $(lhs:1).4b[3]
      sdot $(acc:15).4s, $(rhs:1).16b, $(lhs:1).4b[3]
    )ASM";
  return kernel;
}
```