78 lines · python
1# RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s2 3from mlir.dialects.linalg.opdsl.lang import *4 5 6# CHECK: ---7# CHECK-LABEL: matmul8# CHECK: args:9# CHECK: name: A10# CHECK: kind: input_tensor11# CHECK: type_var: T12# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s0, s1)>13# CHECK: name: B14# CHECK: kind: input_tensor15# CHECK: type_var: T16# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s1, s2)>17# CHECK: name: C18# CHECK: kind: output_tensor19# CHECK: type_var: U20# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>21# CHECK: name: bfn22# CHECK: kind: binary_fn_attr23# CHECK: default_fn: mul24# CHECK: name: ufn25# CHECK: kind: unary_fn_attr26# CHECK: default_fn: exp27# CHECK: name: cast28# CHECK: kind: type_fn_attr29# CHECK: default_fn: cast_signed30@linalg_structured_op31def matmul(32 A=TensorDef(T, S.M, S.K),33 B=TensorDef(T, S.K, S.N),34 C=TensorDef(U, S.M, S.N, output=True),35 bfn=BinaryFnAttrDef(default=BinaryFn.mul),36 ufn=UnaryFnAttrDef(default=UnaryFn.exp),37 cast=TypeFnAttrDef(default=TypeFn.cast_signed),38):39 C[D.m, D.n] += bfn(cast(U, A[D.m, D.k]), cast(U, B[D.k, D.n]))40 41 42# CHECK: ---43# CHECK-LABEL: fill44# CHECK: args:45# CHECK: name: value46# CHECK: kind: scalar47# CHECK-NOT: shape_map:48# CHECK: type_var: T49@linalg_structured_op50def fill(value=ScalarDef(T), O=TensorDef(T, S.M, S.K, output=True)):51 O[D.m, D.n] = value52 53 54# CHECK: ---55# CHECK-LABEL: strided_copy56# CHECK: args:57# CHECK: name: I58# CHECK: kind: input_tensor59# CHECK: type_var: T60# CHECK: shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s0, s1)>61# CHECK: name: O62# CHECK: kind: output_tensor63# CHECK: type_var: T64# CHECK: shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s2, s3)>65# CHECK: name: strides66# CHECK: kind: index_attr67# CHECK: index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s4, s5)>68# CHECK: default_indices:69# CHECK: - 170# CHECK: - 271@linalg_structured_op72def strided_copy(73 I=TensorDef(T, S.IH, S.IW),74 O=TensorDef(T, S.OH, S.OW, output=True),75 strides=IndexAttrDef(S.SH, S.SW, default=[1, 2]),76):77 O[D.oh, D.ow] = I[D.oh * S.SH, D.ow * S.SW]78