1// RUN: mlir-opt %s -test-sparsification="ptr-type=1 ind-type=1" | \ 2// RUN: FileCheck %s --check-prefix=CHECK-TYPE0 3// RUN: mlir-opt %s -test-sparsification="ptr-type=1 ind-type=2" | \ 4// RUN: FileCheck %s --check-prefix=CHECK-TYPE1 5// RUN: mlir-opt %s -test-sparsification="ptr-type=2 ind-type=1" | \ 6// RUN: FileCheck %s --check-prefix=CHECK-TYPE2 7// RUN: mlir-opt %s -test-sparsification="ptr-type=2 ind-type=2" | \ 8// RUN: FileCheck %s --check-prefix=CHECK-TYPE3 9 10#trait_mul_1d = { 11 indexing_maps = [ 12 affine_map<(i) -> (i)>, // a 13 affine_map<(i) -> (i)>, // b 14 affine_map<(i) -> (i)> // x (out) 15 ], 16 sparse = [ 17 [ "S" ], // a 18 [ "D" ], // b 19 [ "D" ] // x 20 ], 21 iterator_types = ["parallel"], 22 doc = "x(i) = a(i) * b(i)" 23} 24 25// CHECK-TYPE0-LABEL: func @mul_dd( 26// CHECK-TYPE0: %[[C0:.*]] = constant 0 : index 27// CHECK-TYPE0: %[[C1:.*]] = constant 1 : index 28// CHECK-TYPE0: %[[P0:.*]] = load %{{.*}}[%[[C0]]] : memref<?xi64> 29// CHECK-TYPE0: %[[B0:.*]] = index_cast %[[P0]] : i64 to index 30// CHECK-TYPE0: %[[P1:.*]] = load %{{.*}}[%[[C1]]] : memref<?xi64> 31// CHECK-TYPE0: %[[B1:.*]] = index_cast %[[P1]] : i64 to index 32// CHECK-TYPE0: scf.for %[[I:.*]] = %[[B0]] to %[[B1]] step %[[C1]] { 33// CHECK-TYPE0: %[[IND0:.*]] = load %{{.*}}[%[[I]]] : memref<?xi64> 34// CHECK-TYPE0: %[[INDC:.*]] = index_cast %[[IND0]] : i64 to index 35// CHECK-TYPE0: %[[VAL0:.*]] = load %{{.*}}[%[[I]]] : memref<?xf64> 36// CHECK-TYPE0: %[[VAL1:.*]] = load %{{.*}}[%[[INDC]]] : memref<32xf64> 37// CHECK-TYPE0: %[[MUL:.*]] = mulf %[[VAL0]], %[[VAL1]] : f64 38// CHECK-TYPE0: store %[[MUL]], %{{.*}}[%[[INDC]]] : memref<32xf64> 39// CHECK-TYPE0: } 40 41// CHECK-TYPE1-LABEL: func @mul_dd( 42// CHECK-TYPE1: %[[C0:.*]] = constant 0 : index 43// CHECK-TYPE1: %[[C1:.*]] = constant 1 : index 44// CHECK-TYPE1: %[[P0:.*]] = load %{{.*}}[%[[C0]]] : memref<?xi64> 45// CHECK-TYPE1: %[[B0:.*]] = index_cast %[[P0]] : i64 to index 46// CHECK-TYPE1: %[[P1:.*]] = load %{{.*}}[%[[C1]]] : memref<?xi64> 47// CHECK-TYPE1: %[[B1:.*]] = index_cast %[[P1]] : i64 to index 48// CHECK-TYPE1: scf.for %[[I:.*]] = %[[B0]] to %[[B1]] step %[[C1]] { 49// CHECK-TYPE1: %[[IND0:.*]] = load %{{.*}}[%[[I]]] : memref<?xi32> 50// CHECK-TYPE1: %[[INDC:.*]] = index_cast %[[IND0]] : i32 to index 51// CHECK-TYPE1: %[[VAL0:.*]] = load %{{.*}}[%[[I]]] : memref<?xf64> 52// CHECK-TYPE1: %[[VAL1:.*]] = load %{{.*}}[%[[INDC]]] : memref<32xf64> 53// CHECK-TYPE1: %[[MUL:.*]] = mulf %[[VAL0]], %[[VAL1]] : f64 54// CHECK-TYPE1: store %[[MUL]], %{{.*}}[%[[INDC]]] : memref<32xf64> 55// CHECK-TYPE1: } 56 57// CHECK-TYPE2-LABEL: func @mul_dd( 58// CHECK-TYPE2: %[[C0:.*]] = constant 0 : index 59// CHECK-TYPE2: %[[C1:.*]] = constant 1 : index 60// CHECK-TYPE2: %[[P0:.*]] = load %{{.*}}[%[[C0]]] : memref<?xi32> 61// CHECK-TYPE2: %[[B0:.*]] = index_cast %[[P0]] : i32 to index 62// CHECK-TYPE2: %[[P1:.*]] = load %{{.*}}[%[[C1]]] : memref<?xi32> 63// CHECK-TYPE2: %[[B1:.*]] = index_cast %[[P1]] : i32 to index 64// CHECK-TYPE2: scf.for %[[I:.*]] = %[[B0]] to %[[B1]] step %[[C1]] { 65// CHECK-TYPE2: %[[IND0:.*]] = load %{{.*}}[%[[I]]] : memref<?xi64> 66// CHECK-TYPE2: %[[INDC:.*]] = index_cast %[[IND0]] : i64 to index 67// CHECK-TYPE2: %[[VAL0:.*]] = load %{{.*}}[%[[I]]] : memref<?xf64> 68// CHECK-TYPE2: %[[VAL1:.*]] = load %{{.*}}[%[[INDC]]] : memref<32xf64> 69// CHECK-TYPE2: %[[MUL:.*]] = mulf %[[VAL0]], %[[VAL1]] : f64 70// CHECK-TYPE2: store %[[MUL]], %{{.*}}[%[[INDC]]] : memref<32xf64> 71// CHECK-TYPE2: } 72 73// CHECK-TYPE3-LABEL: func @mul_dd( 74// CHECK-TYPE3: %[[C0:.*]] = constant 0 : index 75// CHECK-TYPE3: %[[C1:.*]] = constant 1 : index 76// CHECK-TYPE3: %[[P0:.*]] = load %{{.*}}[%[[C0]]] : memref<?xi32> 77// CHECK-TYPE3: %[[B0:.*]] = index_cast %[[P0]] : i32 to index 78// CHECK-TYPE3: %[[P1:.*]] = load %{{.*}}[%[[C1]]] : memref<?xi32> 79// CHECK-TYPE3: %[[B1:.*]] = index_cast %[[P1]] : i32 to index 80// CHECK-TYPE3: scf.for %[[I:.*]] = %[[B0]] to %[[B1]] step %[[C1]] { 81// CHECK-TYPE3: %[[IND0:.*]] = load %{{.*}}[%[[I]]] : memref<?xi32> 82// CHECK-TYPE3: %[[INDC:.*]] = index_cast %[[IND0]] : i32 to index 83// CHECK-TYPE3: %[[VAL0:.*]] = load %{{.*}}[%[[I]]] : memref<?xf64> 84// CHECK-TYPE3: %[[VAL1:.*]] = load %{{.*}}[%[[INDC]]] : memref<32xf64> 85// CHECK-TYPE3: %[[MUL:.*]] = mulf %[[VAL0]], %[[VAL1]] : f64 86// CHECK-TYPE3: store %[[MUL]], %{{.*}}[%[[INDC]]] : memref<32xf64> 87// CHECK-TYPE3: } 88 89func @mul_dd(%arga: tensor<32xf64>, %argb: tensor<32xf64>) -> tensor<32xf64> { 90 %0 = linalg.generic #trait_mul_1d 91 ins(%arga, %argb: tensor<32xf64>, tensor<32xf64>) { 92 ^bb(%a: f64, %b: f64): 93 %0 = mulf %a, %b : f64 94 linalg.yield %0 : f64 95 } -> tensor<32xf64> 96 return %0 : tensor<32xf64> 97} 98 99