This is an archive of the discontinued LLVM Phabricator instance.

Differential D111565

[mlir][Linalg] Enable vectorization of 'mul', 'and', 'or' and 'xor' reductions
ClosedPublic

Authored by dcaballe on Oct 11 2021, 10:48 AM.

Details

Reviewers
pifon2a
ThomasRaoux
nicolasvasilache
herhut
aartbik
Commits
rGeeb09fd6465e: [mlir][Linalg] Enable vectorization of 'mul', 'and', 'or' and 'xor' reductions
Summary

This patch adds support for vectorizing 'mul', 'and', 'or' anx 'xor' reductions
to Linalg.

Diff Detail

Event Timeline

dcaballe created this revision.Oct 11 2021, 10:48 AM
Herald added subscribers: wenzhicui, wrengr, Chia-hungDuan and 18 others. · View Herald TranscriptOct 11 2021, 10:48 AM
dcaballe requested review of this revision.Oct 11 2021, 10:48 AM
Herald added a project: Restricted Project. · View Herald TranscriptOct 11 2021, 10:48 AM
Herald added subscribers: limo1996, stephenneuendorffer. · View Herald Transcript
Harbormaster completed remote builds in B128157: Diff 378734.Oct 11 2021, 10:48 AM
aartbik added inline comments.Oct 11 2021, 10:58 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
127

since you are adding and/or, how about adding xor too?

pifon2a accepted this revision.Oct 11 2021, 11:05 AM
This revision is now accepted and ready to land.Oct 11 2021, 11:05 AM
ThomasRaoux accepted this revision.Oct 11 2021, 12:23 PM
dcaballe updated this revision to Diff 378852.Oct 11 2021, 6:56 PM

Added support for 'xor'

Harbormaster completed remote builds in B128248: Diff 378852.Oct 11 2021, 6:56 PM
dcaballe retitled this revision from [mlir][Linalg] Enable vectorization of 'mul', 'and' and 'or' reductions to [mlir][Linalg] Enable vectorization of 'mul', 'and', 'or' and 'xor' reductions.Oct 11 2021, 6:57 PM
dcaballe edited the summary of this revision. (Show Details)
aartbik accepted this revision.Oct 12 2021, 9:28 AM
pifon2a accepted this revision.Oct 12 2021, 1:33 PM

I think this is ready to be landed.

This revision was landed with ongoing or failed builds.Oct 12 2021, 2:10 PM
Closed by commit rGeeb09fd6465e: [mlir][Linalg] Enable vectorization of 'mul', 'and', 'or' and 'xor' reductions (authored by dcaballe). · Explain Why
This revision was automatically updated to reflect the committed changes.
dcaballe added a commit: rGeeb09fd6465e: [mlir][Linalg] Enable vectorization of 'mul', 'and', 'or' and 'xor' reductions.

Revision Contents

PathSize
mlir/
lib/
Dialect/
Linalg/
Transforms/
3 lines
test/
Dialect/
Linalg/
79 lines

Diff 378734

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Show First 20 Lines • Show All 112 Lines • ▼ Show 20 Lines
static llvm::Optional<vector::CombiningKind> static llvm::Optional<vector::CombiningKind>
getKindForOp(Operation *reductionOp) { getKindForOp(Operation *reductionOp) {
if (!reductionOp) if (!reductionOp)
return llvm::None; return llvm::None;
return llvm::TypeSwitch<Operation *, llvm::Optional<vector::CombiningKind>>( return llvm::TypeSwitch<Operation *, llvm::Optional<vector::CombiningKind>>(
reductionOp) reductionOp)
.Case<AddIOp, AddFOp>([&](auto op) { return vector::CombiningKind::ADD; }) .Case<AddIOp, AddFOp>([&](auto op) { return vector::CombiningKind::ADD; })
.Case<AndOp>([&](auto op) { return vector::CombiningKind::AND; })
.Case<MaxSIOp>([&](auto op) { return vector::CombiningKind::MAXSI; }) .Case<MaxSIOp>([&](auto op) { return vector::CombiningKind::MAXSI; })
.Case<MaxFOp>([&](auto op) { return vector::CombiningKind::MAXF; }) .Case<MaxFOp>([&](auto op) { return vector::CombiningKind::MAXF; })
.Case<MinSIOp>([&](auto op) { return vector::CombiningKind::MINSI; }) .Case<MinSIOp>([&](auto op) { return vector::CombiningKind::MINSI; })
.Case<MinFOp>([&](auto op) { return vector::CombiningKind::MINF; }) .Case<MinFOp>([&](auto op) { return vector::CombiningKind::MINF; })
.Case<MulIOp, MulFOp>([&](auto op) { return vector::CombiningKind::MUL; })
.Case<OrOp>([&](auto op) { return vector::CombiningKind::OR; })
aartbikUnsubmitted
Not Done
ReplyInline Actions

since you are adding and/or, how about adding xor too?

aartbik: since you are adding and/or, how about adding xor too?
.Default([&](auto op) { return llvm::None; }); .Default([&](auto op) { return llvm::None; });
} }
/// Check whether `outputOperand` is a reduction with a single combiner /// Check whether `outputOperand` is a reduction with a single combiner
/// operation. Return the combiner operation kind of the reduction, if /// operation. Return the combiner operation kind of the reduction, if
/// supported. Return llvm::None, otherwise. Multiple reduction operations would /// supported. Return llvm::None, otherwise. Multiple reduction operations would
/// impose an ordering between reduction dimensions and is currently unsupported /// impose an ordering between reduction dimensions and is currently unsupported
/// in Linalg. This limitation is motivated by the fact that e.g. min(max(X)) != /// in Linalg. This limitation is motivated by the fact that e.g. min(max(X)) !=
▲ Show 20 Lines • Show All 1,222 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/vectorization.mlir

Show First 20 Lines • Show All 812 Lines • ▼ Show 20 Lines
func @red_max_2d(%arg0: tensor<4x4xf32>) -> tensor<4xf32> { func @red_max_2d(%arg0: tensor<4x4xf32>) -> tensor<4xf32> {
// CHECK: linalg.init_tensor [4] : tensor<4xf32> // CHECK: linalg.init_tensor [4] : tensor<4xf32>
// CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32>
// CHECK: vector.transfer_read {{.*}} : tensor<4x4xf32>, vector<4x4xf32> // CHECK: vector.transfer_read {{.*}} : tensor<4x4xf32>, vector<4x4xf32>
// CHECK: vector.transfer_read {{.*}} : tensor<4xf32>, vector<4x4xf32> // CHECK: vector.transfer_read {{.*}} : tensor<4xf32>, vector<4x4xf32>
// CHECK: maxf {{.*}} : vector<4x4xf32> // CHECK: maxf {{.*}} : vector<4x4xf32>
// CHECK: vector.multi_reduction #vector.kind<maxf>, {{.*}} [1] : vector<4x4xf32> to vector<4xf32> // CHECK: vector.multi_reduction #vector.kind<maxf>, {{.*}} [1] : vector<4x4xf32> to vector<4xf32>
// CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<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> %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)>, %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) { ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) {
^bb0(%in0: f32, %out0: f32): // no predecessors ^bb0(%in0: f32, %out0: f32): // no predecessors
%max = maxf %in0, %out0 : f32 %max = maxf %in0, %out0 : f32
linalg.yield %max : f32 linalg.yield %max : f32
} -> tensor<4xf32> } -> tensor<4xf32>
Show All 22 Lines ^bb0(%in0: f32, %out0: f32): // no predecessors
%min = minf %in0, %out0 : f32 %min = minf %in0, %out0 : f32
linalg.yield %min : f32 linalg.yield %min : f32
} -> tensor<4xf32> } -> tensor<4xf32>
return %red : tensor<4xf32> return %red : tensor<4xf32>
} }
// ----- // -----
// 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<mul>, {{.*}} [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<or>, {{.*}} [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<and>, {{.*}} [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-DAG: #[[$M5:.*]] = affine_map<(d0, d1) -> (d0, 0)>
// CHECK-LABEL: func @explicit_broadcast( // CHECK-LABEL: func @explicit_broadcast(
func @explicit_broadcast(%arg0: tensor<4x4xf32>, %arg1: tensor<4x1xf32>) -> tensor<4x4xf32> { func @explicit_broadcast(%arg0: tensor<4x4xf32>, %arg1: tensor<4x1xf32>) -> tensor<4x4xf32> {
// CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<4x4xf32>, vector<4x4xf32> // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<4x4xf32>, vector<4x4xf32>
// CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true], permutation_map = #[[$M5]]} : tensor<4x1xf32>, vector<4x4xf32> // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true], permutation_map = #[[$M5]]} : tensor<4x1xf32>, vector<4x4xf32>
// CHECK: subf {{.*}} : vector<4x4xf32> // CHECK: subf {{.*}} : vector<4x4xf32>
// CHECK: vector.transfer_write {{.*}} {in_bounds = [true, true]} : vector<4x4xf32>, tensor<4x4xf32> // CHECK: vector.transfer_write {{.*}} {in_bounds = [true, true]} : vector<4x4xf32>, tensor<4x4xf32>
▲ Show 20 Lines • Show All 48 Lines • Show Last 20 Lines