[all-commits] [llvm/llvm-project] 2ceedc: [mlir][linalg] Add symbolic type conversion to lin...

[Stella Laurenzo via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 2ceedc3a201386c6cbbcea5cec3f5e01d04f6445
      https://github.com/llvm/llvm-project/commit/2ceedc3a201386c6cbbcea5cec3f5e01d04f6445
  Author: Stella Laurenzo <stellaraccident at gmail.com>
  Date:   2021-02-27 (Sat, 27 Feb 2021)

  Changed paths:
    M mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
    M mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
    M mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir
    M mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

  Log Message:
  -----------
  [mlir][linalg] Add symbolic type conversion to linalg named ops.

This enables this kind of construct in the DSL to generate a named op that is polymorphic over numeric type variables `T` and `U`, generating the correct arithmetic casts at construction time:

```
@tc_def_op
def polymorphic_matmul(A=TensorDef(T1, S.M, S.K),
                       B=TensorDef(T2, S.K, S.N),
                       C=TensorDef(U, S.M, S.N, output=True)):
  implements(ContractionOpInterface)
  C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
```

Presently, this only supports type variables that are bound to the element type of one of the arguments, although a further extension that allows binding a type variable to an attribute would allow some more expressiveness and may be useful for some formulations. This is left to a future patch. In addition, this patch does not yet materialize the verifier support which ensures that types are bound correctly (for such simple examples, failing to do so will yield IR that fails verification, it just won't yet fail with a precise error).

Note that the full grid of extensions/truncation/int<->float conversions are supported, but many of them are lossy and higher level code needs to be mindful of numerics (it is not the job of this level).

As-is, this should be sufficient for most integer matmul scenarios we work with in typical quantization schemes.

Differential Revision: https://reviews.llvm.org/D97603