Skip to content

[mlir][spirv] Handle non-innerprod float vector add reductions #73476

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Nov 27, 2023
Merged

[mlir][spirv] Handle non-innerprod float vector add reductions #73476

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
49041e0
[mlir][spirv] Handle non-innerprod float vector add reductions
kuhar Nov 27, 2023
07fa057
Fix typo
kuhar Nov 27, 2023
dce4cdb
Handle one element and accumulator
kuhar Nov 27, 2023
882aa17
Address comments
kuhar Nov 27, 2023
ac69293
Simplify
kuhar Nov 27, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 36 additions & 8 deletions mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
#include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Location.h"
Expand Down Expand Up @@ -755,16 +756,43 @@ struct VectorReductionToFPDotProd final
if (!resultType)
return rewriter.notifyMatchFailure(op, "result is not a float");

auto mul = adaptor.getVector().getDefiningOp<arith::MulFOp>();
if (!mul)
return rewriter.notifyMatchFailure(
op, "reduction operand is not 'arith.mulf'");
Value vec = adaptor.getVector();
Value acc = adaptor.getAcc();

auto vectorType = dyn_cast<VectorType>(vec.getType());
if (!vectorType) {
assert(isa<FloatType>(vec.getType()) &&
"Expected the vector to be scalarized");
if (acc) {
rewriter.replaceOpWithNewOp<spirv::FAddOp>(op, acc, vec);
return success();
}

rewriter.replaceOp(op, vec);
return success();
}

Location loc = op.getLoc();
Value res = rewriter.create<spirv::DotOp>(loc, resultType, mul.getLhs(),
mul.getRhs());
if (op.getAcc())
res = rewriter.create<spirv::FAddOp>(loc, adaptor.getAcc(), res);
Value lhs;
Value rhs;
if (auto mul = vec.getDefiningOp<arith::MulFOp>()) {
lhs = mul.getLhs();
rhs = mul.getRhs();
} else {
// If the operand is not a mul, use a vector of ones for the dot operand
// to just sum up all values.
lhs = vec;
Attribute oneAttr =
rewriter.getFloatAttr(vectorType.getElementType(), 1.0);
oneAttr = SplatElementsAttr::get(vectorType, oneAttr);
rhs = rewriter.create<spirv::ConstantOp>(loc, vectorType, oneAttr);
}
assert(lhs);
assert(rhs);

Value res = rewriter.create<spirv::DotOp>(loc, resultType, lhs, rhs);
if (acc)
res = rewriter.create<spirv::FAddOp>(loc, acc, res);

rewriter.replaceOp(op, res);
return success();
Expand Down
56 changes: 52 additions & 4 deletions mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -500,31 +500,79 @@ func.func @reduction_add(%v : vector<4xi32>) -> i32 {

// -----

// CHECK-LABEL: func @reduction_addf
// CHECK-LABEL: func @reduction_addf_mulf
// CHECK-SAME: (%[[ARG0:.+]]: vector<4xf32>, %[[ARG1:.+]]: vector<4xf32>)
// CHECK: %[[DOT:.+]] = spirv.Dot %[[ARG0]], %[[ARG1]] : vector<4xf32> -> f32
// CHECK: return %[[DOT]] : f32
func.func @reduction_addf(%arg0: vector<4xf32>, %arg1: vector<4xf32>) -> f32 {
func.func @reduction_addf_mulf(%arg0: vector<4xf32>, %arg1: vector<4xf32>) -> f32 {
%mul = arith.mulf %arg0, %arg1 : vector<4xf32>
%red = vector.reduction <add>, %mul : vector<4xf32> into f32
return %red : f32
}

// -----

// CHECK-LABEL: func @reduction_addf_acc
// CHECK-LABEL: func @reduction_addf_acc_mulf
// CHECK-SAME: (%[[ARG0:.+]]: vector<4xf32>, %[[ARG1:.+]]: vector<4xf32>, %[[ACC:.+]]: f32)
// CHECK: %[[DOT:.+]] = spirv.Dot %[[ARG0]], %[[ARG1]] : vector<4xf32> -> f32
// CHECK: %[[RES:.+]] = spirv.FAdd %[[ACC]], %[[DOT]] : f32
// CHECK: return %[[RES]] : f32
func.func @reduction_addf_acc(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %acc: f32) -> f32 {
func.func @reduction_addf_acc_mulf(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %acc: f32) -> f32 {
%mul = arith.mulf %arg0, %arg1 : vector<4xf32>
%red = vector.reduction <add>, %mul, %acc : vector<4xf32> into f32
return %red : f32
}

// -----

// CHECK-LABEL: func @reduction_addf
// CHECK-SAME: (%[[ARG0:.+]]: vector<4xf32>)
// CHECK: %[[ONE:.+]] = spirv.Constant dense<1.0{{.+}}> : vector<4xf32>
// CHECK: %[[DOT:.+]] = spirv.Dot %[[ARG0]], %[[ONE]] : vector<4xf32> -> f32
// CHECK: return %[[DOT]] : f32
func.func @reduction_addf_mulf(%arg0: vector<4xf32>) -> f32 {
%red = vector.reduction <add>, %arg0 : vector<4xf32> into f32
return %red : f32
}

// -----

// CHECK-LABEL: func @reduction_addf_acc
// CHECK-SAME: (%[[ARG0:.+]]: vector<4xf32>, %[[ACC:.+]]: f32)
// CHECK: %[[ONE:.+]] = spirv.Constant dense<1.0{{.*}}> : vector<4xf32>
// CHECK: %[[DOT:.+]] = spirv.Dot %[[ARG0]], %[[ONE]] : vector<4xf32> -> f32
// CHECK: %[[RES:.+]] = spirv.FAdd %[[ACC]], %[[DOT]] : f32
// CHECK: return %[[RES]] : f32
func.func @reduction_addf_acc(%arg0: vector<4xf32>, %acc: f32) -> f32 {
%red = vector.reduction <add>, %arg0, %acc : vector<4xf32> into f32
return %red : f32
}

// -----

// CHECK-LABEL: func @reduction_addf_one_elem
// CHECK-SAME: (%[[ARG0:.+]]: vector<1xf32>)
// CHECK: %[[RES:.+]] = builtin.unrealized_conversion_cast %[[ARG0]] : vector<1xf32> to f32
// CHECK: return %[[RES]] : f32
func.func @reduction_addf_one_elem(%arg0: vector<1xf32>) -> f32 {
%red = vector.reduction <add>, %arg0 : vector<1xf32> into f32
return %red : f32
}

// -----

// CHECK-LABEL: func @reduction_addf_one_elem_acc
// CHECK-SAME: (%[[ARG0:.+]]: vector<1xf32>, %[[ACC:.+]]: f32)
// CHECK: %[[RHS:.+]] = builtin.unrealized_conversion_cast %[[ARG0]] : vector<1xf32> to f32
// CHECK: %[[RES:.+]] = spirv.FAdd %[[ACC]], %[[RHS]] : f32
// CHECK: return %[[RES]] : f32
func.func @reduction_addf_one_elem_acc(%arg0: vector<1xf32>, %acc: f32) -> f32 {
%red = vector.reduction <add>, %arg0, %acc : vector<1xf32> into f32
return %red : f32
}

// -----

// CHECK-LABEL: func @reduction_mul
// CHECK-SAME: (%[[V:.+]]: vector<3xf32>, %[[S:.+]]: f32)
// CHECK: %[[S0:.+]] = spirv.CompositeExtract %[[V]][0 : i32] : vector<3xf32>
Expand Down