1// RUN: mlir-opt %s -sparsification="vectorization-strategy=0 vl=16" | \ 2// RUN: FileCheck %s --check-prefix=CHECK-VEC0 3// RUN: mlir-opt %s -sparsification="vectorization-strategy=1 vl=16" | \ 4// RUN: FileCheck %s --check-prefix=CHECK-VEC1 5// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=16" | \ 6// RUN: FileCheck %s --check-prefix=CHECK-VEC2 7 8#DenseVector = #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }> 9 10#SparseVector = #sparse_tensor.encoding<{ 11 dimLevelType = [ "compressed" ], 12 pointerBitWidth = 32, 13 indexBitWidth = 32 14}> 15 16#SparseMatrix = #sparse_tensor.encoding<{ 17 dimLevelType = [ "dense", "compressed" ], 18 pointerBitWidth = 32, 19 indexBitWidth = 32 20}> 21 22#trait_scale_d = { 23 indexing_maps = [ 24 affine_map<(i) -> (i)>, // a 25 affine_map<(i) -> (i)> // x (out) 26 ], 27 iterator_types = ["parallel"], 28 doc = "x(i) = a(i) * b" 29} 30 31// 32// CHECK-VEC0-LABEL: func @scale_d 33// CHECK-VEC0-DAG: %[[c0:.*]] = constant 0 : index 34// CHECK-VEC0-DAG: %[[c1:.*]] = constant 1 : index 35// CHECK-VEC0-DAG: %[[c1024:.*]] = constant 1024 : index 36// CHECK-VEC0: scf.for %[[i:.*]] = %[[c0]] to %[[c1024]] step %[[c1]] { 37// CHECK-VEC0: %[[l:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xf32> 38// CHECK-VEC0: %[[m:.*]] = mulf %[[l]], %{{.*}} : f32 39// CHECK-VEC0: store %[[m]], %{{.*}}[%[[i]]] : memref<1024xf32> 40// CHECK-VEC0: } 41// CHECK-VEC0: return 42// 43// CHECK-VEC1-LABEL: func @scale_d 44// CHECK-VEC1-DAG: %[[c0:.*]] = constant 0 : index 45// CHECK-VEC1-DAG: %[[c16:.*]] = constant 16 : index 46// CHECK-VEC1-DAG: %[[c1024:.*]] = constant 1024 : index 47// CHECK-VEC1: scf.for %[[i:.*]] = %[[c0]] to %[[c1024]] step %[[c16]] { 48// CHECK-VEC1: %[[r:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xf32>, vector<16xf32> 49// CHECK-VEC1: %[[b:.*]] = vector.broadcast %{{.*}} : f32 to vector<16xf32> 50// CHECK-VEC1: %[[m:.*]] = mulf %[[r]], %[[b]] : vector<16xf32> 51// CHECK-VEC1: vector.store %[[m]], %{{.*}}[%[[i]]] : memref<1024xf32>, vector<16xf32> 52// CHECK-VEC1: } 53// CHECK-VEC1: return 54// 55// CHECK-VEC2-LABEL: func @scale_d 56// CHECK-VEC2-DAG: %[[c0:.*]] = constant 0 : index 57// CHECK-VEC2-DAG: %[[c16:.*]] = constant 16 : index 58// CHECK-VEC2-DAG: %[[c1024:.*]] = constant 1024 : index 59// CHECK-VEC2: scf.for %[[i:.*]] = %[[c0]] to %[[c1024]] step %[[c16]] { 60// CHECK-VEC2: %[[r:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xf32>, vector<16xf32> 61// CHECK-VEC2: %[[b:.*]] = vector.broadcast %{{.*}} : f32 to vector<16xf32> 62// CHECK-VEC2: %[[m:.*]] = mulf %[[r]], %[[b]] : vector<16xf32> 63// CHECK-VEC2: vector.store %[[m]], %{{.*}}[%[[i]]] : memref<1024xf32>, vector<16xf32> 64// CHECK-VEC2: } 65// CHECK-VEC2: return 66// 67func @scale_d(%arga: tensor<1024xf32, #DenseVector>, %b: f32, %argx: tensor<1024xf32>) -> tensor<1024xf32> { 68 %0 = linalg.generic #trait_scale_d 69 ins(%arga: tensor<1024xf32, #DenseVector>) 70 outs(%argx: tensor<1024xf32>) { 71 ^bb(%a: f32, %x: f32): 72 %0 = mulf %a, %b : f32 73 linalg.yield %0 : f32 74 } -> tensor<1024xf32> 75 return %0 : tensor<1024xf32> 76} 77 78#trait_mul_s = { 79 indexing_maps = [ 80 affine_map<(i) -> (i)>, // a 81 affine_map<(i) -> (i)>, // b 82 affine_map<(i) -> (i)> // x (out) 83 ], 84 iterator_types = ["parallel"], 85 doc = "x(i) = a(i) * b(i)" 86} 87 88// 89// CHECK-VEC0-LABEL: func @mul_s 90// CHECK-VEC0-DAG: %[[c0:.*]] = constant 0 : index 91// CHECK-VEC0-DAG: %[[c1:.*]] = constant 1 : index 92// CHECK-VEC0: %[[p:.*]] = memref.load %{{.*}}[%[[c0]]] : memref<?xi32> 93// CHECK-VEC0: %[[a:.*]] = zexti %[[p]] : i32 to i64 94// CHECK-VEC0: %[[q:.*]] = index_cast %[[a]] : i64 to index 95// CHECK-VEC0: %[[r:.*]] = memref.load %{{.*}}[%[[c1]]] : memref<?xi32> 96// CHECK-VEC0: %[[b:.*]] = zexti %[[r]] : i32 to i64 97// CHECK-VEC0: %[[s:.*]] = index_cast %[[b]] : i64 to index 98// CHECK-VEC0: scf.for %[[i:.*]] = %[[q]] to %[[s]] step %[[c1]] { 99// CHECK-VEC0: %[[li:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xi32> 100// CHECK-VEC0: %[[zi:.*]] = zexti %[[li]] : i32 to i64 101// CHECK-VEC0: %[[ci:.*]] = index_cast %[[zi]] : i64 to index 102// CHECK-VEC0: %[[la:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xf32> 103// CHECK-VEC0: %[[lb:.*]] = memref.load %{{.*}}[%[[ci]]] : memref<1024xf32> 104// CHECK-VEC0: %[[m:.*]] = mulf %[[la]], %[[lb]] : f32 105// CHECK-VEC0: store %[[m]], %{{.*}}[%[[ci]]] : memref<1024xf32> 106// CHECK-VEC0: } 107// CHECK-VEC0: return 108// 109// CHECK-VEC1-LABEL: func @mul_s 110// CHECK-VEC1-DAG: %[[c0:.*]] = constant 0 : index 111// CHECK-VEC1-DAG: %[[c1:.*]] = constant 1 : index 112// CHECK-VEC1: %[[p:.*]] = memref.load %{{.*}}[%[[c0]]] : memref<?xi32> 113// CHECK-VEC1: %[[a:.*]] = zexti %[[p]] : i32 to i64 114// CHECK-VEC1: %[[q:.*]] = index_cast %[[a]] : i64 to index 115// CHECK-VEC1: %[[r:.*]] = memref.load %{{.*}}[%[[c1]]] : memref<?xi32> 116// CHECK-VEC1: %[[b:.*]] = zexti %[[r]] : i32 to i64 117// CHECK-VEC1: %[[s:.*]] = index_cast %[[b]] : i64 to index 118// CHECK-VEC1: scf.for %[[i:.*]] = %[[q]] to %[[s]] step %[[c1]] { 119// CHECK-VEC1: %[[li:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xi32> 120// CHECK-VEC1: %[[zi:.*]] = zexti %[[li]] : i32 to i64 121// CHECK-VEC1: %[[ci:.*]] = index_cast %[[zi]] : i64 to index 122// CHECK-VEC1: %[[la:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xf32> 123// CHECK-VEC1: %[[lb:.*]] = memref.load %{{.*}}[%[[ci]]] : memref<1024xf32> 124// CHECK-VEC1: %[[m:.*]] = mulf %[[la]], %[[lb]] : f32 125// CHECK-VEC1: store %[[m]], %{{.*}}[%[[ci]]] : memref<1024xf32> 126// CHECK-VEC1: } 127// CHECK-VEC1: return 128// 129// CHECK-VEC2-LABEL: func @mul_s 130// CHECK-VEC2-DAG: %[[c0:.*]] = constant 0 : index 131// CHECK-VEC2-DAG: %[[c1:.*]] = constant 1 : index 132// CHECK-VEC2-DAG: %[[c16:.*]] = constant 16 : index 133// CHECK-VEC2: %[[p:.*]] = memref.load %{{.*}}[%[[c0]]] : memref<?xi32> 134// CHECK-VEC2: %[[a:.*]] = zexti %[[p]] : i32 to i64 135// CHECK-VEC2: %[[q:.*]] = index_cast %[[a]] : i64 to index 136// CHECK-VEC2: %[[r:.*]] = memref.load %{{.*}}[%[[c1]]] : memref<?xi32> 137// CHECK-VEC2: %[[b:.*]] = zexti %[[r]] : i32 to i64 138// CHECK-VEC2: %[[s:.*]] = index_cast %[[b]] : i64 to index 139// CHECK-VEC2: scf.for %[[i:.*]] = %[[q]] to %[[s]] step %[[c16]] { 140// CHECK-VEC2: %[[sub:.*]] = subi %[[s]], %[[i]] : index 141// CHECK-VEC2: %[[mask:.*]] = vector.create_mask %[[sub]] : vector<16xi1> 142// CHECK-VEC2: %[[li:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %{{.*}} : memref<?xi32>, vector<16xi1>, vector<16xi32> into vector<16xi32> 143// CHECK-VEC2: %[[zi:.*]] = zexti %[[li]] : vector<16xi32> to vector<16xi64> 144// CHECK-VEC2: %[[la:.*]] = vector.maskedload %{{.*}}[%[[i]]], %[[mask]], %{{.*}} : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32> 145// CHECK-VEC2: %[[lb:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[zi]]], %[[mask]], %{{.*}} : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32> 146// CHECK-VEC2: %[[m:.*]] = mulf %[[la]], %[[lb]] : vector<16xf32> 147// CHECK-VEC2: vector.scatter %{{.*}}[%[[c0]]] [%[[zi]]], %[[mask]], %[[m]] : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> 148// CHECK-VEC2: } 149// CHECK-VEC2: return 150// 151func @mul_s(%arga: tensor<1024xf32, #SparseVector>, %argb: tensor<1024xf32>, %argx: tensor<1024xf32>) -> tensor<1024xf32> { 152 %0 = linalg.generic #trait_mul_s 153 ins(%arga, %argb: tensor<1024xf32, #SparseVector>, tensor<1024xf32>) 154 outs(%argx: tensor<1024xf32>) { 155 ^bb(%a: f32, %b: f32, %x: f32): 156 %0 = mulf %a, %b : f32 157 linalg.yield %0 : f32 158 } -> tensor<1024xf32> 159 return %0 : tensor<1024xf32> 160} 161 162#trait_reduction_d = { 163 indexing_maps = [ 164 affine_map<(i) -> (i)>, // a 165 affine_map<(i) -> (i)>, // b 166 affine_map<(i) -> ()> // x (out) 167 ], 168 iterator_types = ["reduction"], 169 doc = "x += a(i) * b(i)" 170} 171 172// 173// CHECK-VEC0-LABEL: func @reduction_d 174// CHECK-VEC0-DAG: %[[c0:.*]] = constant 0 : index 175// CHECK-VEC0-DAG: %[[c1:.*]] = constant 1 : index 176// CHECK-VEC0-DAG: %[[c1024:.*]] = constant 1024 : index 177// CHECK-VEC0: %[[red:.*]] = scf.for %[[i:.*]] = %[[c0]] to %[[c1024]] step %[[c1]] iter_args(%[[red_in:.*]] = %{{.*}}) -> (f32) { 178// CHECK-VEC0: %[[la:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xf32> 179// CHECK-VEC0: %[[lb:.*]] = memref.load %{{.*}}[%[[i]]] : memref<1024xf32> 180// CHECK-VEC0: %[[m:.*]] = mulf %[[la]], %[[lb]] : f32 181// CHECK-VEC0: %[[a:.*]] = addf %[[red_in]], %[[m]] : f32 182// CHECK-VEC0: scf.yield %[[a]] : f32 183// CHECK-VEC0: } 184// CHECK-VEC0: return 185// 186// CHECK-VEC1-LABEL: func @reduction_d 187// CHECK-VEC1-DAG: %[[c0:.*]] = constant 0 : index 188// CHECK-VEC1-DAG: %[[c16:.*]] = constant 16 : index 189// CHECK-VEC1-DAG: %[[c1024:.*]] = constant 1024 : index 190// CHECK-VEC1-DAG: %[[v0:.*]] = constant dense<0.000000e+00> : vector<16xf32> 191// CHECK-VEC1: %[[red:.*]] = scf.for %[[i:.*]] = %[[c0]] to %[[c1024]] step %[[c16]] iter_args(%[[red_in:.*]] = %[[v0]]) -> (vector<16xf32>) { 192// CHECK-VEC1: %[[la:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xf32>, vector<16xf32> 193// CHECK-VEC1: %[[lb:.*]] = vector.load %{{.*}}[%[[i]]] : memref<1024xf32>, vector<16xf32> 194// CHECK-VEC1: %[[m:.*]] = mulf %[[la]], %[[lb]] : vector<16xf32> 195// CHECK-VEC1: %[[a:.*]] = addf %[[red_in]], %[[m]] : vector<16xf32> 196// CHECK-VEC1: scf.yield %[[a]] : vector<16xf32> 197// CHECK-VEC1: } 198// CHECK-VEC1: %{{.*}} = vector.reduction "add", %[[red]], %{{.*}} : vector<16xf32> into f32 199// CHECK-VEC1: return 200// 201// CHECK-VEC2-LABEL: func @reduction_d 202// CHECK-VEC2-DAG: %[[c0:.*]] = constant 0 : index 203// CHECK-VEC2-DAG: %[[c16:.*]] = constant 16 : index 204// CHECK-VEC2-DAG: %[[c1024:.*]] = constant 1024 : index 205// CHECK-VEC2-DAG: %[[v0:.*]] = constant dense<0.000000e+00> : vector<16xf32> 206// CHECK-VEC2: %[[red:.*]] = scf.for %[[i:.*]] = %[[c0]] to %[[c1024]] step %[[c16]] iter_args(%[[red_in:.*]] = %[[v0]]) -> (vector<16xf32>) { 207// CHECK-VEC2: %[[la:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xf32>, vector<16xf32> 208// CHECK-VEC2: %[[lb:.*]] = vector.load %{{.*}}[%[[i]]] : memref<1024xf32>, vector<16xf32> 209// CHECK-VEC2: %[[m:.*]] = mulf %[[la]], %[[lb]] : vector<16xf32> 210// CHECK-VEC2: %[[a:.*]] = addf %[[red_in]], %[[m]] : vector<16xf32> 211// CHECK-VEC2: scf.yield %[[a]] : vector<16xf32> 212// CHECK-VEC2: } 213// CHECK-VEC2: %{{.*}} = vector.reduction "add", %[[red]], %{{.*}} : vector<16xf32> into f32 214// CHECK-VEC2: return 215// 216func @reduction_d(%arga: tensor<1024xf32, #DenseVector>, %argb: tensor<1024xf32>, %argx: tensor<f32>) -> tensor<f32> { 217 %0 = linalg.generic #trait_reduction_d 218 ins(%arga, %argb: tensor<1024xf32, #DenseVector>, tensor<1024xf32>) 219 outs(%argx: tensor<f32>) { 220 ^bb(%a: f32, %b: f32, %x: f32): 221 %0 = mulf %a, %b : f32 222 %1 = addf %x, %0 : f32 223 linalg.yield %1 : f32 224 } -> tensor<f32> 225 return %0 : tensor<f32> 226} 227 228#trait_mul_ds = { 229 indexing_maps = [ 230 affine_map<(i,j) -> (i,j)>, // A 231 affine_map<(i,j) -> (i,j)>, // B 232 affine_map<(i,j) -> (i,j)> // X (out) 233 ], 234 iterator_types = ["parallel", "parallel"], 235 doc = "X(i,j) = A(i,j) * B(i,j)" 236} 237 238// 239// CHECK-VEC0-LABEL: func @mul_ds 240// CHECK-VEC0-DAG: %[[c0:.*]] = constant 0 : index 241// CHECK-VEC0-DAG: %[[c1:.*]] = constant 1 : index 242// CHECK-VEC0-DAG: %[[c512:.*]] = constant 512 : index 243// CHECK-VEC0: scf.for %[[i:.*]] = %[[c0]] to %[[c512]] step %[[c1]] { 244// CHECK-VEC0: %[[p:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xi32> 245// CHECK-VEC0: %[[a:.*]] = zexti %[[p]] : i32 to i64 246// CHECK-VEC0: %[[q:.*]] = index_cast %[[a]] : i64 to index 247// CHECK-VEC0: %[[a:.*]] = addi %[[i]], %[[c1]] : index 248// CHECK-VEC0: %[[r:.*]] = memref.load %{{.*}}[%[[a]]] : memref<?xi32> 249// CHECK-VEC0: %[[b:.*]] = zexti %[[r]] : i32 to i64 250// CHECK-VEC0: %[[s:.*]] = index_cast %[[b]] : i64 to index 251// CHECK-VEC0: scf.for %[[j:.*]] = %[[q]] to %[[s]] step %[[c1]] { 252// CHECK-VEC0: %[[lj:.*]] = memref.load %{{.*}}[%[[j]]] : memref<?xi32> 253// CHECK-VEC0: %[[zj:.*]] = zexti %[[lj]] : i32 to i64 254// CHECK-VEC0: %[[cj:.*]] = index_cast %[[zj]] : i64 to index 255// CHECK-VEC0: %[[la:.*]] = memref.load %{{.*}}[%[[j]]] : memref<?xf32> 256// CHECK-VEC0: %[[lb:.*]] = memref.load %{{.*}}[%[[i]], %[[cj]]] : memref<512x1024xf32> 257// CHECK-VEC0: %[[m:.*]] = mulf %[[la]], %[[lb]] : f32 258// CHECK-VEC0: store %[[m]], %{{.*}}[%[[i]], %[[cj]]] : memref<512x1024xf32> 259// CHECK-VEC0: } 260// CHECK-VEC0: } 261// CHECK-VEC0: return 262// 263// CHECK-VEC1-LABEL: func @mul_ds 264// CHECK-VEC1-DAG: %[[c0:.*]] = constant 0 : index 265// CHECK-VEC1-DAG: %[[c1:.*]] = constant 1 : index 266// CHECK-VEC1-DAG: %[[c512:.*]] = constant 512 : index 267// CHECK-VEC1: scf.for %[[i:.*]] = %[[c0]] to %[[c512]] step %[[c1]] { 268// CHECK-VEC1: %[[p:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xi32> 269// CHECK-VEC1: %[[a:.*]] = zexti %[[p]] : i32 to i64 270// CHECK-VEC1: %[[q:.*]] = index_cast %[[a]] : i64 to index 271// CHECK-VEC1: %[[a:.*]] = addi %[[i]], %[[c1]] : index 272// CHECK-VEC1: %[[r:.*]] = memref.load %{{.*}}[%[[a]]] : memref<?xi32> 273// CHECK-VEC1: %[[b:.*]] = zexti %[[r]] : i32 to i64 274// CHECK-VEC1: %[[s:.*]] = index_cast %[[b]] : i64 to index 275// CHECK-VEC1: scf.for %[[j:.*]] = %[[q]] to %[[s]] step %[[c1]] { 276// CHECK-VEC1: %[[lj:.*]] = memref.load %{{.*}}[%[[j]]] : memref<?xi32> 277// CHECK-VEC1: %[[zj:.*]] = zexti %[[lj]] : i32 to i64 278// CHECK-VEC1: %[[cj:.*]] = index_cast %[[zj]] : i64 to index 279// CHECK-VEC1: %[[la:.*]] = memref.load %{{.*}}[%[[j]]] : memref<?xf32> 280// CHECK-VEC1: %[[lb:.*]] = memref.load %{{.*}}[%[[i]], %[[cj]]] : memref<512x1024xf32> 281// CHECK-VEC1: %[[m:.*]] = mulf %[[la]], %[[lb]] : f32 282// CHECK-VEC1: store %[[m]], %{{.*}}[%[[i]], %[[cj]]] : memref<512x1024xf32> 283// CHECK-VEC1: } 284// CHECK-VEC1: } 285// CHECK-VEC1: return 286// 287// CHECK-VEC2-LABEL: func @mul_ds 288// CHECK-VEC2-DAG: %[[c0:.*]] = constant 0 : index 289// CHECK-VEC2-DAG: %[[c1:.*]] = constant 1 : index 290// CHECK-VEC2-DAG: %[[c16:.*]] = constant 16 : index 291// CHECK-VEC2-DAG: %[[c512:.*]] = constant 512 : index 292// CHECK-VEC2: scf.for %[[i:.*]] = %[[c0]] to %[[c512]] step %[[c1]] { 293// CHECK-VEC2: %[[p:.*]] = memref.load %{{.*}}[%[[i]]] : memref<?xi32> 294// CHECK-VEC2: %[[a:.*]] = zexti %[[p]] : i32 to i64 295// CHECK-VEC2: %[[q:.*]] = index_cast %[[a]] : i64 to index 296// CHECK-VEC2: %[[a:.*]] = addi %[[i]], %[[c1]] : index 297// CHECK-VEC2: %[[r:.*]] = memref.load %{{.*}}[%[[a]]] : memref<?xi32> 298// CHECK-VEC2: %[[b:.*]] = zexti %[[r]] : i32 to i64 299// CHECK-VEC2: %[[s:.*]] = index_cast %[[b]] : i64 to index 300// CHECK-VEC2: scf.for %[[j:.*]] = %[[q]] to %[[s]] step %[[c16]] { 301// CHECK-VEC2: %[[sub:.*]] = subi %[[s]], %[[j]] : index 302// CHECK-VEC2: %[[mask:.*]] = vector.create_mask %[[sub]] : vector<16xi1> 303// CHECK-VEC2: %[[lj:.*]] = vector.maskedload %{{.*}}[%[[j]]], %[[mask]], %{{.*}} : memref<?xi32>, vector<16xi1>, vector<16xi32> into vector<16xi32> 304// CHECK-VEC2: %[[zj:.*]] = zexti %[[lj]] : vector<16xi32> to vector<16xi64> 305// CHECK-VEC2: %[[la:.*]] = vector.maskedload %{{.*}}[%[[j]]], %[[mask]], %{{.*}} : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32> 306// CHECK-VEC2: %[[lb:.*]] = vector.gather %{{.*}}[%[[i]], %[[c0]]] [%[[zj]]], %[[mask]], %{{.*}} : memref<512x1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32> 307// CHECK-VEC2: %[[m:.*]] = mulf %[[la]], %[[lb]] : vector<16xf32> 308// CHECK-VEC2: vector.scatter %{{.*}}[%[[i]], %[[c0]]] [%[[zj]]], %[[mask]], %[[m]] : memref<512x1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> 309// CHECK-VEC2: } 310// CHECK-VEC2: } 311// CHECK-VEC2: return 312// 313func @mul_ds(%arga: tensor<512x1024xf32, #SparseMatrix>, %argb: tensor<512x1024xf32>, %argx: tensor<512x1024xf32>) -> tensor<512x1024xf32> { 314 %0 = linalg.generic #trait_mul_ds 315 ins(%arga, %argb: tensor<512x1024xf32, #SparseMatrix>, tensor<512x1024xf32>) 316 outs(%argx: tensor<512x1024xf32>) { 317 ^bb(%a: f32, %b: f32, %x: f32): 318 %0 = mulf %a, %b : f32 319 linalg.yield %0 : f32 320 } -> tensor<512x1024xf32> 321 return %0 : tensor<512x1024xf32> 322} 323