116 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: assignments:9# CHECK: -10# CHECK: arg: C11# CHECK: value:12# CHECK: scalar_fn:13# CHECK: kind: binary14# CHECK: fn_name: add15# CHECK: operands:16# CHECK: scalar_fn:17# CHECK: kind: binary18# CHECK: attr_name: mul19# CHECK: operands:20# CHECK: scalar_fn:21# CHECK: kind: type22# CHECK: attr_name: cast23# CHECK: type_var: U24# CHECK: operands:25# CHECK: scalar_arg: A26# CHECK: scalar_fn:27# CHECK: kind: type28# CHECK: attr_name: cast29# CHECK: type_var: U30# CHECK: operands:31# CHECK: scalar_arg: B32@linalg_structured_op33def matmul(34 A=TensorDef(T, S.M, S.K),35 B=TensorDef(T, S.K, S.N),36 C=TensorDef(U, S.M, S.N, output=True),37 mul=BinaryFnAttrDef(default=BinaryFn.mul),38 cast=TypeFnAttrDef(default=TypeFn.cast_signed),39):40 C[D.m, D.n] += mul(cast(U, A[D.m, D.k]), cast(U, B[D.k, D.n]))41 42 43# CHECK: ---44# CHECK-LABEL: constants45# CHECK: assignments:46# CHECK: -47# CHECK: arg: O48# CHECK: scalar_fn:49# CHECK: kind: binary50# CHECK: fn_name: sub51# CHECK: operands:52# CHECK: scalar_fn:53# CHECK: kind: binary54# CHECK: fn_name: add55# CHECK: operands:56# CHECK: scalar_fn:57# CHECK: kind: unary58# CHECK: fn_name: exp59# CHECK: operands:60# CHECK: scalar_fn:61# CHECK: kind: type62# CHECK: type_var: T63# CHECK: operands:64# CHECK: scalar_const: '3.1415926535897931 : f64'65# CHECK: scalar_fn:66# CHECK: kind: type67# CHECK: fn_name: cast_signed68# CHECK: type_var: T69# CHECK: operands:70# CHECK: scalar_const: '42 : i64'71# CHECK: scalar_fn:72# CHECK: kind: type73# CHECK: fn_name: cast_signed74# CHECK: type_var: T75# CHECK: operands:76# CHECK: scalar_fn:77# CHECK: kind: unary78# CHECK: attr_name: exp79# CHECK: operands:80# CHECK: scalar_const: '1.{{[0]*}}e+03 : f64'81@linalg_structured_op82def constants(83 O=TensorDef(T, S.M, S.K, output=True), exp=UnaryFnAttrDef(default=UnaryFn.exp)84):85 pi = TypeFn.cast_signed(T, const(3.1415926535897931))86 cst42 = TypeFn.cast_signed(T, const(42))87 cst1000 = TypeFn.cast_signed(T, exp(const(1e3)))88 O[D.m, D.n] = UnaryFn.exp(pi) + cst42 - cst100089 90 91# CHECK: ---92# CHECK-LABEL: indices93# CHECK: assignments:94# CHECK: -95# CHECK: arg: O96# CHECK: scalar_fn:97# CHECK: kind: binary98# CHECK: fn_name: add99# CHECK: operands:100# CHECK: scalar_index: 1101# CHECK: scalar_index: 0102@linalg_structured_op103def indices(O=TensorDef(T, S.M, S.K, output=True)):104 O[D.m, D.n] = index(D.n) + index(D.m)105 106 107# CHECK: ---108# CHECK-LABEL: fill109# CHECK: assignments:110# CHECK: -111# CHECK: arg: O112# CHECK: scalar_arg: value113@linalg_structured_op114def fill(value=ScalarDef(T), O=TensorDef(T, S.M, S.K, output=True)):115 O[D.m, D.n] = value116