1// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=16" -for-loop-peeling -canonicalize -split-input-file | \
2// RUN:   FileCheck %s
3
4#SparseVector = #sparse_tensor.encoding<{
5  dimLevelType = [ "compressed" ],
6  pointerBitWidth = 32,
7  indexBitWidth = 32
8}>
9
10#trait_mul_s = {
11  indexing_maps = [
12    affine_map<(i) -> (i)>,  // a
13    affine_map<(i) -> (i)>,  // b
14    affine_map<(i) -> (i)>   // x (out)
15  ],
16  iterator_types = ["parallel"],
17  doc = "x(i) = a(i) * b(i)"
18}
19
20// CHECK-DAG:   #[[$map0:.*]] = affine_map<()[s0, s1] -> (s0 + ((-s0 + s1) floordiv 16) * 16)>
21// CHECK-DAG:   #[[$map1:.*]] = affine_map<()[s0, s1] -> ((s0 - s1) mod 16)>
22// CHECK-LABEL: func @mul_s
23// CHECK-DAG:   %[[c0:.*]] = constant 0 : index
24// CHECK-DAG:   %[[c1:.*]] = constant 1 : index
25// CHECK-DAG:   %[[c16:.*]] = constant 16 : index
26// CHECK:       %[[p:.*]] = memref.load %{{.*}}[%[[c0]]] : memref<?xi32>
27// CHECK:       %[[a:.*]] = zexti %[[p]] : i32 to i64
28// CHECK:       %[[q:.*]] = index_cast %[[a]] : i64 to index
29// CHECK:       %[[r:.*]] = memref.load %{{.*}}[%[[c1]]] : memref<?xi32>
30// CHECK:       %[[b:.*]] = zexti %[[r]] : i32 to i64
31// CHECK:       %[[s:.*]] = index_cast %[[b]] : i64 to index
32// CHECK:       %[[boundary:.*]] = affine.apply #[[$map0]]()[%[[q]], %[[s]]]
33// CHECK:       scf.for %[[i:.*]] = %[[q]] to %[[boundary]] step %[[c16]] {
34// CHECK:         %[[mask:.*]] = vector.constant_mask [16] : vector<16xi1>
35// CHECK:         %[[li:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xi32>, vector<16xi32>
36// CHECK:         %[[zi:.*]] = zexti %[[li]] : vector<16xi32> to vector<16xi64>
37// CHECK:         %[[la:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xf32>, vector<16xf32>
38// CHECK:         %[[lb:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[zi]]], %[[mask]], %{{.*}} : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32>
39// CHECK:         %[[m:.*]] = mulf %[[la]], %[[lb]] : vector<16xf32>
40// CHECK:         vector.scatter %{{.*}}[%[[c0]]] [%[[zi]]], %[[mask]], %[[m]] : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32>
41// CHECK:       }
42// CHECK:       %[[has_more:.*]] = cmpi slt, %[[boundary]], %[[s]] : index
43// CHECK:       scf.if %[[has_more]] {
44// CHECK:         %[[sub:.*]] = affine.apply #[[$map1]]()[%[[s]], %[[q]]]
45// CHECK:         %[[mask2:.*]] = vector.create_mask %[[sub]] : vector<16xi1>
46// CHECK:         %[[li2:.*]] = vector.maskedload %{{.*}}[%[[boundary]]], %[[mask2]], %{{.*}} : memref<?xi32>, vector<16xi1>, vector<16xi32> into vector<16xi32>
47// CHECK:         %[[zi2:.*]] = zexti %[[li2]] : vector<16xi32> to vector<16xi64>
48// CHECK:         %[[la2:.*]] = vector.maskedload %{{.*}}[%[[boundary]]], %[[mask2]], %{{.*}} : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
49// CHECK:         %[[lb2:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[zi2]]], %[[mask2]], %{{.*}} : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32>
50// CHECK:         %[[m2:.*]] = mulf %[[la2]], %[[lb2]] : vector<16xf32>
51// CHECK:         vector.scatter %{{.*}}[%[[c0]]] [%[[zi2]]], %[[mask2]], %[[m2]] : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32>
52// CHECK:       }
53// CHECK:       return
54//
55func @mul_s(%arga: tensor<1024xf32, #SparseVector>, %argb: tensor<1024xf32>, %argx: tensor<1024xf32>) -> tensor<1024xf32> {
56  %0 = linalg.generic #trait_mul_s
57    ins(%arga, %argb: tensor<1024xf32, #SparseVector>, tensor<1024xf32>)
58    outs(%argx: tensor<1024xf32>) {
59      ^bb(%a: f32, %b: f32, %x: f32):
60        %0 = mulf %a, %b : f32
61        linalg.yield %0 : f32
62  } -> tensor<1024xf32>
63  return %0 : tensor<1024xf32>
64}
65