Hi everyone,
I am proposing adding support to declaratively create semantically charged “named ops” in Linalg.
The end goal is that we will be able to automatically generate a large fraction of ops in a dialects from a declarative specification that resembles (e.g. for linalg.matmul):
"C(i,j) +=! A(i,k) * B(k,j) where i in 0:M, k in 0:K, j in 0:N"
Concretely, the plan is to start with a specific Tablegen backend to automatically populate the attributes of linalg.generic from a textual specification. As things progress and evolve, refactorings are expected to take place and be reusable across dialects (e.g. more idiomatic support for affine maps in Tablegen).
Such a mechanism is expected to help build a system to understand tradeoffs related to :
- succinct declaration of many ops from different application domains, ML in particular
- source of mathematical truth that carries the semantics of the op at a high-level
- op “class” vs op “instance” specification and rank-polymorphism
- roundtripping between the “named” form and the “generic” form of an op and implications on matchers and library calls / CISC ISAs
- transformations that apply to the “named” form (e.g. TASO/TensorRT/Algebraic rewrites) or the “generic” form (fusion, tiling etc) of the op.
- shape inference
- specification language design
Progress on these fronts will be on a per-need basis and with input from the community, but at least we can start from something concrete.
While the backend would like to be unopinionated on the choice of the syntax, it will need to start somewhere, so I am proposing to use a syntax inspired by TensorComprehensions to get started.
Concretely, the first revision will assume some parser exists and add the ability to write:
def PointwiseAddOp : LinalgNamedStructured_Op<"pointwise_add", [
NInputs<2>,
NOutputs<1>,
NamedStructuredOpTraits]> {
let arguments = (ins Variadic<LinalgOperand>:$views);
let results = (outs Variadic<AnyRankedTensor>:$output_tensors);
let spec = "A(...) = B(...) + C(...)";
let assemblyFormat = "`(` operands `)` attr-dict `:` "
"functional-type(operands, results)";
}
def PointwiseMulOp : LinalgNamedStructured_Op<"pointwise_mul", [
NInputs<2>,
NOutputs<1>,
NamedStructuredOpTraits]> {
let arguments = (ins Variadic<LinalgOperand>:$views);
let results = (outs Variadic<AnyRankedTensor>:$output_tensors);
let spec = "A(...) = B(...) * C(...)";
let assemblyFormat = "`(` operands `)` attr-dict `:` "
"functional-type(operands, results)";
}
The spec strings will just be checked for equality for now, a parser will be developed later.
This will result in the following legal IR on arbitrarily mixed tensors and buffers:
func @pointwise1d(%a: memref<?xf32>, %b: memref<?xf32>, %c: memref<?xf32>,
%ta: tensor<?xf32>, %tb: tensor<?xf32>, %tc: tensor<?xf32>) {
linalg.pointwise_add(%a, %b, %c): (memref<?xf32>, memref<?xf32>, memref<?xf32>) -> ()
%dadd = linalg.pointwise_add(%a, %b): (memref<?xf32>, memref<?xf32>) -> tensor<?xf32>
%eadd = linalg.pointwise_add(%ta, %b): (tensor<?xf32>, memref<?xf32>) -> tensor<?xf32>
linalg.pointwise_add(%ta, %tb, %c): (tensor<?xf32>, tensor<?xf32>, memref<?xf32>) -> ()
linalg.pointwise_mul(%a, %b, %c): (memref<?xf32>, memref<?xf32>, memref<?xf32>) -> ()
%dmul = linalg.pointwise_mul(%a, %b): (memref<?xf32>, memref<?xf32>) -> tensor<?xf32>
%emul = linalg.pointwise_mul(%ta, %b): (tensor<?xf32>, memref<?xf32>) -> tensor<?xf32>
linalg.pointwise_mul(%ta, %tb, %c): (tensor<?xf32>, tensor<?xf32>, memref<?xf32>) -> ()
return
}
Followup revisions will consist in (1) enabling the roundtripping between named and generic form as well as (2) lowering to loops.
This is also related to the discussion on TCP ops and tensors.
I will send an RFC revision soon so people can comment on the approach.
Please let me know what you think, thanks!
) it’s easy to be very agile with it and evolve it.