73 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# Verify that simple case with iteration order defined lexically and reduction7# dims auto discovered emits the right shape, indexing maps and iterator types.8# CHECK: ---9# CHECK-LABEL: matmul10# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s0, s1)>11# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s1, s2)>12# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>13# CHECK: static_indexing_maps:14# CHECK-NEXT: - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>15# CHECK-NEXT: - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2, d1)>16# CHECK-NEXT: - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>17# CHECK: iterator_types:18# CHECK-NEXT: - parallel19# CHECK-NEXT: - parallel20# CHECK-NEXT: - reduction21@linalg_structured_op22def matmul(23 A=TensorDef(T, S.M, S.K),24 B=TensorDef(T, S.K, S.N),25 C=TensorDef(U, S.M, S.N, output=True),26):27 domain(D.m, D.n, D.k)28 C[D.m, D.n] += TypeFn.cast_signed(U, A[D.m, D.k]) * TypeFn.cast_signed(29 U, B[D.k, D.n]30 )31 32 33# Verifies that assignment to a scalar (represented as [None]) is represented34# correctly.35# CHECK: ---36# CHECK-LABEL: dot37# CHECK: shape_map: affine_map<()[s0] -> (s0)>38# CHECK: shape_map: affine_map<()[s0] -> (s0)>39# CHECK: shape_map: affine_map<()[s0] -> ()>40# CHECK: static_indexing_maps:41# CHECK-NEXT: - affine_map<(d0)[s0] -> (d0)>42# CHECK-NEXT: - affine_map<(d0)[s0] -> (d0)>43# CHECK-NEXT: - affine_map<(d0)[s0] -> ()>44# CHECK: iterator_types:45# CHECK-NEXT: - reduction46@linalg_structured_op47def dot(A=TensorDef(T, S.M), B=TensorDef(T, S.M), C=TensorDef(U, output=True)):48 C[None] += TypeFn.cast_signed(U, A[D.m]) * TypeFn.cast_signed(U, B[D.m])49 50 51# Verifies that the index_dims of shape-only operands translate to correct52# indexing maps.53# CHECK: ---54# CHECK-LABEL: pool55# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s0)>56# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s1)>57# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s2)>58# CHECK: static_indexing_maps:59# CHECK-NEXT: - affine_map<(d0, d1)[s0, s1, s2] -> (d0 * 2 + d1)>60# CHECK-NEXT: - affine_map<(d0, d1)[s0, s1, s2] -> (d1)>61# CHECK-NEXT: - affine_map<(d0, d1)[s0, s1, s2] -> (d0)>62# CHECK: iterator_types:63# CHECK-NEXT: - parallel64# CHECK-NEXT: - reduction65@linalg_structured_op66def pool(67 I=TensorDef(T, S.I),68 K=TensorDef(T, S.K, index_dims=[D.k]),69 O=TensorDef(U, S.O, output=True),70):71 domain(D.o, D.k)72 O[D.o] += TypeFn.cast_signed(U, I[D.o * 2 + D.k])73