1// RUN: mlir-opt %s --sparse-reinterpret-map --sparsification -cse -sparse-vectorization="vl=16" -scf-for-loop-peeling -canonicalize -cse | \ 2// RUN: FileCheck %s 3 4#SparseVector = #sparse_tensor.encoding<{ 5 map = (d0) -> (d0 : compressed), 6 posWidth = 32, 7 crdWidth = 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<(d0)[s0] -> (-d0 + s0)> 22// CHECK-LABEL: func @mul_s 23// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index 24// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index 25// CHECK-DAG: %[[c16:.*]] = arith.constant 16 : index 26// CHECK: %[[p:.*]] = memref.load %{{.*}}[%[[c0]]] : memref<?xi32> 27// CHECK: %[[a:.*]] = arith.extui %[[p]] : i32 to i64 28// CHECK: %[[q:.*]] = arith.index_cast %[[a]] : i64 to index 29// CHECK: %[[r:.*]] = memref.load %{{.*}}[%[[c1]]] : memref<?xi32> 30// CHECK: %[[b:.*]] = arith.extui %[[r]] : i32 to i64 31// CHECK: %[[s:.*]] = arith.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:.*]] = arith.extui %[[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:.*]] = arith.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: scf.for %[[i2:.*]] = %[[boundary]] to %[[s]] step %[[c16]] { 43// CHECK: %[[sub:.*]] = affine.apply #[[$map1]](%[[i2]])[%[[s]]] 44// CHECK: %[[mask2:.*]] = vector.create_mask %[[sub]] : vector<16xi1> 45// CHECK: %[[li2:.*]] = vector.maskedload %{{.*}}[%[[i2]]], %[[mask2]], %{{.*}} : memref<?xi32>, vector<16xi1>, vector<16xi32> into vector<16xi32> 46// CHECK: %[[zi2:.*]] = arith.extui %[[li2]] : vector<16xi32> to vector<16xi64> 47// CHECK: %[[la2:.*]] = vector.maskedload %{{.*}}[%[[i2]]], %[[mask2]], %{{.*}} : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32> 48// CHECK: %[[lb2:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[zi2]]], %[[mask2]], %{{.*}} : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32> 49// CHECK: %[[m2:.*]] = arith.mulf %[[la2]], %[[lb2]] : vector<16xf32> 50// CHECK: vector.scatter %{{.*}}[%[[c0]]] [%[[zi2]]], %[[mask2]], %[[m2]] : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> 51// CHECK: } 52// CHECK: return 53// 54func.func @mul_s(%arga: tensor<1024xf32, #SparseVector>, %argb: tensor<1024xf32>, %argx: tensor<1024xf32>) -> tensor<1024xf32> { 55 %0 = linalg.generic #trait_mul_s 56 ins(%arga, %argb: tensor<1024xf32, #SparseVector>, tensor<1024xf32>) 57 outs(%argx: tensor<1024xf32>) { 58 ^bb(%a: f32, %b: f32, %x: f32): 59 %0 = arith.mulf %a, %b : f32 60 linalg.yield %0 : f32 61 } -> tensor<1024xf32> 62 return %0 : tensor<1024xf32> 63} 64