diff --git a/mlir/docs/Tools/LinalgOpDsl.md b/mlir/docs/Tools/LinalgOpDsl.md --- a/mlir/docs/Tools/LinalgOpDsl.md +++ b/mlir/docs/Tools/LinalgOpDsl.md @@ -54,7 +54,8 @@ them to the same data type as the accumulator/output. """ implements(ContractionOpInterface) - C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n]) + for m, n, k in domain(D.m, D.n, D.k): + C[m, n] += cast(U, A[m, k]) * cast(U, B[k, n]) ``` Here we have a simple type polymorphic contraction that takes arguments `A` @@ -74,32 +75,97 @@ ## Parameters -Structured operations can take two types of parameters namely input/output -tensors and captures. Assignment expressions index the tensor parameters to -access the individual elements, while captures are scalars that can be -accessed directly. +Structured operations take two types of runtime parameters namely scalars and +tensors. While scalars are inputs only, a tensor may be marked as an output. +Assignment expressions index the tensor parameters to access the individual +elements, while scalars can be accessed directly. The following example demonstrates the use of the two parameter types: ```python @linalg_structured_op -def copy_and_scale(I=TensorDef(T, S.M, S.K), - O=TensorDef(T, S.M, S.K, output=True), - val=CaptureDef(T)): - """Scale the input by the captured value and store the result""" - O[D.m, D.n] = I[D.m, D.n] * val +def copy_and_scale(val=ScalarDef(T), + I=TensorDef(T, S.M, S.K), + O=TensorDef(T, S.M, S.K, output=True)): + """Scale the input by the scalar value and store the result""" + for m, n in domain(D.m, D.n): + O[m, n] = I[m, n] * val ``` -The operation scales the input tensor `I` scales its elements by the value -`val` and writes the result to the output tensor `out`. The capture `val` is -bound to a `CaptureDef`, which specifies the type of the captured value. The -tensors are bound to a `TensorDef` as demonstrated by the matmul example. All -parameters appear in the parameter list of the operation: +The operation scales the input tensor `I` scales its elements by the value `val` +and writes the result to the output tensor `out`. The scalar `val` is bound to a +`ScalarDef`, which specifies the type of the scalar operand. The tensors are +bound to a `TensorDef` as demonstrated by the matmul example. All parameters +appear in the parameter list of the operation: ```python -fill(in_tensor, outs=[out_tensor], captures=[captured_val]) +fill(val, in_tensor, outs=[out_tensor]) ``` +## Attributes + +Attributes are compile-time constant parameters only accessible in index +expressions. They can be used to parameterize the access pattern of a structured +operation, for example, by setting its strides. They cannot take part in the +actual computation. + +The following example demonstrates the use of attributes: + +```python +@linalg_structured_op +def strided_copy(I=TensorDef(T, S.IH, S.IW), + O=TensorDef(T, S.OH, S.OW, output=True), + strides=AttributeDef(S.SH, S.SW)): + """Copy a subset of the input tensor elements to the output tensor""" + for oh, ow in domain(D.oh, D.ow): + O[oh, ow] = I[oh * S.SH, ow * S.SW] +``` + +The operation implements a strided copy from the input tensor `I` to the output +tensor `O`. The `strides` attribute is bound to an `AttributeDef`. It defines +the symbols `S.SH` and `S.SW`, which are used to index the input tensor `I`. +When instantiating the operation, the attribute is set using a named argument: + +```python +strided_copy(in_tensor, outs=[out_tensor], strides=[1,2]) +``` + +The `strides` vector elements substitute the symbols `S.SH` and `S.SW` in the +index expressions of the operation instance. + +Attributes are currently limited to integer vectors and only accessible in index +expressions. An operation may have multiple attributes all of them placed at the +end of the parameter list after the output tensors. + +## Shape-Only Tensors + +Structured operations derive the iteration space given the sizes of the input +and output tensors. Certain operations need shape-only tensors that are not +accessed and exist purely for the sake of specifying the iteration domain. An +example is the pooling operation that takes a shape-only tensor to define the +iteration space of the reduction. As shape-only tensors have no uses, the +`TensorDef` takes an additional optional `index_dims` parameter to map the shape +to index dimensions. + +The following example demonstrates the index dimension annotation: + +```python +@linalg_structured_op +def pooling_poly( + I=TensorDef(T1, S.N, S.H, S.W, S.C), + K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), + O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), + strides=AttributeDef(S.SH, S.SW), + dilations=AttributeDef(S.DH, S.DW)): + for n, oh, ow, kh, kw, c in domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c): + O[n, oh, ow, c] += \ + cast(U, I[n, oh * S.SH + kh * S.DH, ow * S.SW + kw * S.DW, c]) +``` + +The pooling operation does not access the shape-only tensor `K`. Instead, the +shapes `S.KH` and `S.KW` specify the iteration domain for the reduction +dimensions `D.kh` and `D.kw`. + ## Assignments The bulk of language consists of assignment expressions of the form above.