4361bd9b7b
Extend the OpDSL syntax with an optional `domain` function to specify an explicit dimension order. The extension is needed to provide more control over the dimension order instead of deducing it implicitly depending on the formulation of the tensor comprehension. Additionally, the patch also ensures the symbols are ordered according to the operand definitions of the operation. Differential Revision: https://reviews.llvm.org/D105117
62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
# RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s
|
|
|
|
from mlir.dialects.linalg.opdsl.lang import *
|
|
|
|
|
|
# CHECK: ---
|
|
# CHECK-LABEL: matmul
|
|
# CHECK: args:
|
|
# CHECK: name: A
|
|
# CHECK: usage: InputOperand
|
|
# CHECK: type_var: T
|
|
# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s0, s1)>
|
|
# CHECK: name: B
|
|
# CHECK: usage: InputOperand
|
|
# CHECK: type_var: T
|
|
# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s1, s2)>
|
|
# CHECK: name: C
|
|
# CHECK: usage: OutputOperand
|
|
# CHECK: type_var: U
|
|
# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>
|
|
@linalg_structured_op
|
|
def matmul(
|
|
A=TensorDef(T, S.M, S.K),
|
|
B=TensorDef(T, S.K, S.N),
|
|
C=TensorDef(U, S.M, S.N, output=True)):
|
|
C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
|
|
|
|
|
|
# CHECK: ---
|
|
# CHECK-LABEL: fill
|
|
# CHECK: args:
|
|
# CHECK: name: value
|
|
# CHECK: usage: InputOperand
|
|
# CHECK-NOT: shape_map:
|
|
# CHECK: type_var: T
|
|
@linalg_structured_op
|
|
def fill(value=ScalarDef(T), O=TensorDef(T, S.M, S.K, output=True)):
|
|
O[D.m, D.n] = value
|
|
|
|
|
|
# CHECK: ---
|
|
# CHECK-LABEL: strided_copy
|
|
# CHECK: args:
|
|
# CHECK: name: I
|
|
# CHECK: usage: InputOperand
|
|
# CHECK: type_var: T
|
|
# CHECK: shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s0, s1)>
|
|
# CHECK: name: O
|
|
# CHECK: usage: OutputOperand
|
|
# CHECK: type_var: T
|
|
# CHECK: shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s2, s3)>
|
|
# CHECK: name: strides
|
|
# CHECK: usage: IndexAttribute
|
|
# CHECK: type_var: I64
|
|
# CHECK: attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s4, s5)>
|
|
@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)):
|
|
O[D.oh, D.ow] = I[D.oh * S.SH, D.ow * S.SW]
|