1# RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s 2 3from mlir.dialects.linalg.opdsl.lang import * 4 5 6# Verify that simple case with iteration order defined lexically and reduction 7# dims auto discovered emits the right shape, indexing maps and iterator types. 8# CHECK: --- 9# CHECK-LABEL: matmul 10# 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: - parallel 19# CHECK-NEXT: - parallel 20# CHECK-NEXT: - reduction 21@linalg_structured_op 22def 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 represented 34# correctly. 35# CHECK: --- 36# CHECK-LABEL: dot 37# 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: - reduction 46@linalg_structured_op 47def 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 correct 52# indexing maps. 53# CHECK: --- 54# CHECK-LABEL: pool 55# 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: - parallel 64# CHECK-NEXT: - reduction 65@linalg_structured_op 66def 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