diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -118,10 +118,13 @@ return llvm::TypeSwitch>( reductionOp) .Case([&](auto op) { return vector::CombiningKind::ADD; }) + .Case([&](auto op) { return vector::CombiningKind::AND; }) .Case([&](auto op) { return vector::CombiningKind::MAXSI; }) .Case([&](auto op) { return vector::CombiningKind::MAXF; }) .Case([&](auto op) { return vector::CombiningKind::MINSI; }) .Case([&](auto op) { return vector::CombiningKind::MINF; }) + .Case([&](auto op) { return vector::CombiningKind::MUL; }) + .Case([&](auto op) { return vector::CombiningKind::OR; }) .Default([&](auto op) { return llvm::None; }); } diff --git a/mlir/test/Dialect/Linalg/vectorization.mlir b/mlir/test/Dialect/Linalg/vectorization.mlir --- a/mlir/test/Dialect/Linalg/vectorization.mlir +++ b/mlir/test/Dialect/Linalg/vectorization.mlir @@ -818,9 +818,9 @@ // CHECK: maxf {{.*}} : vector<4x4xf32> // CHECK: vector.multi_reduction #vector.kind, {{.*}} [1] : vector<4x4xf32> to vector<4xf32> // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> - %minf32 = constant -3.40282e+38 : f32 + %ident = constant -3.40282e+38 : f32 %init = linalg.init_tensor [4] : tensor<4xf32> - %fill = linalg.fill(%minf32, %init) : f32, tensor<4xf32> -> tensor<4xf32> + %fill = linalg.fill(%ident, %init) : f32, tensor<4xf32> -> tensor<4xf32> %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} @@ -859,6 +859,81 @@ // ----- +// CHECK-LABEL: func @red_mul_2d( +func @red_mul_2d(%arg0: tensor<4x4xf32>) -> tensor<4xf32> { + // CHECK: linalg.init_tensor [4] : tensor<4xf32> + // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> + // CHECK: vector.transfer_read {{.*}} : tensor<4x4xf32>, vector<4x4xf32> + // CHECK: vector.transfer_read {{.*}} : tensor<4xf32>, vector<4x4xf32> + // CHECK: mulf {{.*}} : vector<4x4xf32> + // CHECK: vector.multi_reduction #vector.kind, {{.*}} [1] : vector<4x4xf32> to vector<4xf32> + // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> + %ident = constant 1.0 : f32 + %init = linalg.init_tensor [4] : tensor<4xf32> + %fill = linalg.fill(%ident, %init) : f32, tensor<4xf32> -> tensor<4xf32> + %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d0)>], + iterator_types = ["parallel", "reduction"]} + ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) { + ^bb0(%in0: f32, %out0: f32): // no predecessors + %mul = mulf %in0, %out0 : f32 + linalg.yield %mul : f32 + } -> tensor<4xf32> + return %red : tensor<4xf32> +} + +// ----- + +// CHECK-LABEL: func @red_or_2d( +func @red_or_2d(%arg0: tensor<4x4xi1>) -> tensor<4xi1> { + // CHECK: linalg.init_tensor [4] : tensor<4xi1> + // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> + // CHECK: vector.transfer_read {{.*}} : tensor<4x4xi1>, vector<4x4xi1> + // CHECK: vector.transfer_read {{.*}} : tensor<4xi1>, vector<4x4xi1> + // CHECK: or {{.*}} : vector<4x4xi1> + // CHECK: vector.multi_reduction #vector.kind, {{.*}} [1] : vector<4x4xi1> to vector<4xi1> + // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> + %ident = constant false + %init = linalg.init_tensor [4] : tensor<4xi1> + %fill = linalg.fill(%ident, %init) : i1, tensor<4xi1> -> tensor<4xi1> + %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d0)>], + iterator_types = ["parallel", "reduction"]} + ins(%arg0 : tensor<4x4xi1>) outs(%fill : tensor<4xi1>) { + ^bb0(%in0: i1, %out0: i1): // no predecessors + %or = or %in0, %out0 : i1 + linalg.yield %or : i1 + } -> tensor<4xi1> + return %red : tensor<4xi1> +} + +// ----- + +// CHECK-LABEL: func @red_and_2d( +func @red_and_2d(%arg0: tensor<4x4xi1>) -> tensor<4xi1> { + // CHECK: linalg.init_tensor [4] : tensor<4xi1> + // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> + // CHECK: vector.transfer_read {{.*}} : tensor<4x4xi1>, vector<4x4xi1> + // CHECK: vector.transfer_read {{.*}} : tensor<4xi1>, vector<4x4xi1> + // CHECK: and {{.*}} : vector<4x4xi1> + // CHECK: vector.multi_reduction #vector.kind, {{.*}} [1] : vector<4x4xi1> to vector<4xi1> + // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> + %ident = constant true + %init = linalg.init_tensor [4] : tensor<4xi1> + %fill = linalg.fill(%ident, %init) : i1, tensor<4xi1> -> tensor<4xi1> + %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, + affine_map<(d0, d1) -> (d0)>], + iterator_types = ["parallel", "reduction"]} + ins(%arg0 : tensor<4x4xi1>) outs(%fill : tensor<4xi1>) { + ^bb0(%in0: i1, %out0: i1): // no predecessors + %and = and %in0, %out0 : i1 + linalg.yield %and : i1 + } -> tensor<4xi1> + return %red : tensor<4xi1> +} + +// ----- + // CHECK-DAG: #[[$M5:.*]] = affine_map<(d0, d1) -> (d0, 0)> // CHECK-LABEL: func @explicit_broadcast(