Skip to content

Commit b153433

Browse files
d-smirnovd-smirnov
committed
Addressed comments
1 parent e081ff9 commit b153433

File tree

3 files changed

+86
-113
lines changed

3 files changed

+86
-113
lines changed

‎mlir/lib/Dialect/Linalg/Transforms/WinogradConv2D.cpp

Lines changed: 56 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -729,6 +729,7 @@ Value outputTransform(RewriterBase &rewriter, Location loc, Value value,
729729

730730
auto buildBody = [&](OpBuilder &builder, Location loc, ValueRange ivs,
731731
ValueRange args) -> scf::ValueVector {
732+
auto context = builder.getContext();
732733
Value tileHIter = ivs[0];
733734
Value tileWIter = ivs[1];
734735
Value NIter = ivs[2];
@@ -740,29 +741,41 @@ Value outputTransform(RewriterBase &rewriter, Location loc, Value value,
740741
FIter, 2, 3, /*loopNorFIdx=*/4,
741742
/*loopCorFIdx=*/5, /*heightIdx=*/0, /*widthIdx=*/1);
742743

743-
TransformMapKeyTy key = {m, r};
744-
int64_t retRows = 1;
745-
int64_t retCols = 1;
746-
int64_t leftScalarFactor = 1;
747-
int64_t rightScalarFactor = 1;
744+
const TransformMapKeyTy key = {m, r};
745+
const TransformMatrix &AMatrix = AMatrices.at(key);
746+
const TransformMatrix &ATMatrix = ATMatrices.at(key);
747+
int64_t scalarFactor = (rightTransform ? AMatrix.scalarFactor : 1) *
748+
(leftTransform ? ATMatrix.scalarFactor : 1);
749+
int64_t retCols = rightTransform ? AMatrix.cols : 1;
750+
int64_t retRows = leftTransform ? ATMatrix.rows : 1;
751+
748752
Value matmulRetValue = extractValue;
749753
Value zero = builder.create<arith::ConstantOp>(
750754
loc, rewriter.getZeroAttr(elementType));
751-
if (leftTransform) {
752-
// Get constant transform matrix AT.
753-
auto it = ATMatrices.find(key);
754-
if (it == ATMatrices.end())
755-
return {};
756-
const TransformMatrix &ATMatrix = it->second;
757755

758-
leftScalarFactor = ATMatrix.scalarFactor;
759-
retRows = ATMatrix.rows;
756+
auto affineMap =
757+
AffineMap::get(1, 0, {builder.getAffineDimExpr(0) * m}, context);
758+
Value heightOffset =
759+
builder.create<affine::AffineApplyOp>(loc, affineMap, tileHIter);
760+
Value widthOffset =
761+
builder.create<affine::AffineApplyOp>(loc, affineMap, tileWIter);
762+
763+
Value outInitVal =
764+
extract2DDataFrom4D(builder, loc, args[0], NIter, FIter, heightOffset,
765+
widthOffset, retRows, retCols,
766+
/*loopNorFIdx=*/0,
767+
/*loopCorFIdx=*/3, /*heightIdx=*/1,
768+
/*widthIdx=*/2);
769+
if (leftTransform) {
760770
auto matmulType = RankedTensorType::get({retRows, valueW}, elementType);
761-
auto empty =
762-
builder
763-
.create<tensor::EmptyOp>(loc, matmulType.getShape(), elementType)
764-
.getResult();
765-
auto init = builder.create<linalg::FillOp>(loc, zero, empty).getResult(0);
771+
Value init = outInitVal;
772+
if (rightTransform || scalarFactor != 1) {
773+
auto empty = builder
774+
.create<tensor::EmptyOp>(loc, matmulType.getShape(),
775+
elementType)
776+
.getResult();
777+
init = builder.create<linalg::FillOp>(loc, zero, empty).getResult(0);
778+
}
766779

767780
Value AT = create2DTransformMatrix(builder, loc, ATMatrix, elementType);
768781
// Multiply AT x m.
@@ -772,21 +785,16 @@ Value outputTransform(RewriterBase &rewriter, Location loc, Value value,
772785
}
773786

774787
if (rightTransform) {
775-
// Get constant transform matrix T.
776-
auto it = AMatrices.find(key);
777-
if (it == AMatrices.end())
778-
return {};
779-
const TransformMatrix &AMatrix = it->second;
780-
781-
rightScalarFactor = AMatrix.scalarFactor;
782788
auto matmulType =
783789
RankedTensorType::get({retRows, AMatrix.cols}, elementType);
784-
retCols = AMatrix.cols;
785-
auto empty =
786-
builder
787-
.create<tensor::EmptyOp>(loc, matmulType.getShape(), elementType)
788-
.getResult();
789-
auto init = builder.create<linalg::FillOp>(loc, zero, empty).getResult(0);
790+
Value init = outInitVal;
791+
if (scalarFactor != 1) {
792+
auto empty = builder
793+
.create<tensor::EmptyOp>(loc, matmulType.getShape(),
794+
elementType)
795+
.getResult();
796+
init = builder.create<linalg::FillOp>(loc, zero, empty).getResult(0);
797+
}
790798

791799
Value A = create2DTransformMatrix(builder, loc, AMatrix, elementType);
792800
// Multiply y = (AT x m) x A.
@@ -795,66 +803,39 @@ Value outputTransform(RewriterBase &rewriter, Location loc, Value value,
795803
matmulRetValue = matmulOp.getResult(0);
796804
}
797805

798-
if (leftScalarFactor * rightScalarFactor != 1) {
799-
// Multiply scalar factor.
800-
Value scalarFactor = builder.create<arith::ConstantOp>(
801-
loc,
802-
FloatAttr::get(elementType, leftScalarFactor * rightScalarFactor));
806+
if (scalarFactor != 1) {
807+
// Multiply by scalar factor and add outInitVal.
808+
Value scalarFactorValue = builder.create<arith::ConstantOp>(
809+
loc, FloatAttr::get(elementType, scalarFactor));
803810
auto matmulType = RankedTensorType::get({retRows, retCols}, elementType);
804-
auto init = builder.create<tensor::EmptyOp>(loc, matmulType.getShape(),
805-
elementType);
806-
807811
auto identityAffineMap = rewriter.getMultiDimIdentityMap(2);
808812
SmallVector<AffineMap> affineMaps = {
809-
AffineMap::get(2, 0, init.getContext()), identityAffineMap};
810-
auto broadcastedScalar =
813+
AffineMap::get(2, 0, context), identityAffineMap, identityAffineMap};
814+
815+
matmulRetValue =
811816
rewriter
812817
.create<linalg::GenericOp>(
813-
loc, matmulType, ValueRange{scalarFactor}, ValueRange{init},
814-
affineMaps,
818+
loc, matmulType,
819+
ValueRange{scalarFactorValue, matmulRetValue},
820+
ValueRange{outInitVal}, affineMaps,
815821
llvm::ArrayRef<utils::IteratorType>{
816822
utils::IteratorType::parallel,
817823
utils::IteratorType::parallel},
818824
[&](OpBuilder &nestedBuilder, Location nestedLoc,
819825
ValueRange args) {
820-
nestedBuilder.create<linalg::YieldOp>(nestedLoc, args[0]);
826+
auto mulf = nestedBuilder.create<arith::MulFOp>(
827+
nestedLoc, args[0], args[1]);
828+
auto addf = nestedBuilder.create<arith::AddFOp>(
829+
nestedLoc, mulf.getResult(), args[2]);
830+
nestedBuilder.create<linalg::YieldOp>(nestedLoc,
831+
addf.getResult());
821832
})
822833
.getResult(0);
823-
824-
matmulRetValue = builder
825-
.create<linalg::MulOp>(
826-
loc, matmulType,
827-
ValueRange{broadcastedScalar, matmulRetValue},
828-
ValueRange{init})
829-
.getResult(0);
830834
}
831835

832-
auto context = builder.getContext();
833-
auto affineMap =
834-
AffineMap::get(1, 0, {builder.getAffineDimExpr(0) * m}, context);
835-
Value heightOffset =
836-
builder.create<affine::AffineApplyOp>(loc, affineMap, tileHIter);
837-
Value widthOffset =
838-
builder.create<affine::AffineApplyOp>(loc, affineMap, tileWIter);
839-
840-
Value outInitVal =
841-
extract2DDataFrom4D(builder, loc, args[0], NIter, FIter, heightOffset,
842-
widthOffset, retRows, retCols,
843-
/*loopNorFIdx=*/0,
844-
/*loopCorFIdx=*/3, /*heightIdx=*/1,
845-
/*widthIdx=*/2);
846-
Value outVal =
847-
builder
848-
.create<linalg::AddOp>(
849-
loc, outInitVal.getType(), ValueRange{matmulRetValue, outInitVal},
850-
ValueRange{builder.create<tensor::EmptyOp>(
851-
loc, llvm::cast<ShapedType>(outInitVal.getType()).getShape(),
852-
elementType)})
853-
.getResult(0);
854-
855836
// Insert (H, W) to (N, H, W, F).
856837
Value combinedVal =
857-
insert2DDataTo4D(builder, loc, outVal, args[0], NIter, FIter,
838+
insert2DDataTo4D(builder, loc, matmulRetValue, args[0], NIter, FIter,
858839
heightOffset, widthOffset, retRows, retCols,
859840
/*loopNorFIdx=*/0,
860841
/*loopCorFIdx=*/3, /*heightIdx=*/1,

‎mlir/test/Dialect/Linalg/transform-tile-and-winograd-rewrite.mlir

Lines changed: 21 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -97,22 +97,20 @@ module attributes {transform.with_named_sequence} {
9797
// CHECK: %[[S12:.*]] = scf.for %[[ARG7:.*]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG8:.*]] = %[[EXTRACTED_SLICE_7]])
9898
// CHECK: %[[S15:.*]] = scf.for %[[ARG9:.*]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG10:.*]] = %[[ARG8]])
9999
// CHECK: %[[EXTRACTED_SLICE_8:.*]] = tensor.extract_slice %[[EXTRACTED_SLICE]][0, 0, 0, 0, %[[ARG7]], %[[ARG9]]] [6, 6, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1]
100+
// CHECK: %[[S25:.*]] = tensor.extract_slice %[[ARG10]][%[[ARG7]], 0, 0, %[[ARG9]]] [1, 4, 4, 1] [1, 1, 1, 1]
100101
// CHECK: %[[S16:.*]] = tensor.empty() : tensor<4x6xf32>
101102
// CHECK: %[[S17:.*]] = linalg.fill ins(%[[CST_6]] : f32) outs(%[[S16]] : tensor<4x6xf32>) -> tensor<4x6xf32>
102103
// CHECK: %[[S18:.*]] = linalg.matmul ins(%[[CST_1]], %[[EXTRACTED_SLICE_8]] : tensor<4x6xf32>, tensor<6x6xf32>) outs(%[[S17]] : tensor<4x6xf32>) -> tensor<4x6xf32>
103104
// CHECK: %[[S19:.*]] = tensor.empty() : tensor<4x4xf32>
104105
// CHECK: %[[S20:.*]] = linalg.fill ins(%[[CST_6]] : f32) outs(%[[S19]] : tensor<4x4xf32>) -> tensor<4x4xf32>
105106
// CHECK: %[[S21:.*]] = linalg.matmul ins(%[[S18]], %[[CST_0]] : tensor<4x6xf32>, tensor<6x4xf32>) outs(%[[S20]] : tensor<4x4xf32>) -> tensor<4x4xf32>
106-
// CHECK: %[[S22:.*]] = tensor.empty() : tensor<4x4xf32>
107-
// CHECK: %[[S23:.*]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%[[CST]] : f32) outs(%[[S22]] : tensor<4x4xf32>) {
108-
// CHECK: ^bb0(%[[IN:.*]]: f32, %[[OUT:.*]]: f32):
109-
// CHECK: linalg.yield %[[IN]] : f32
107+
// CHECK: %[[S23:.*]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%[[CST]], %[[S21]] : f32, tensor<4x4xf32>) outs(%[[S25]] : tensor<4x4xf32>) {
108+
// CHECK: ^bb0(%[[IN1:.*]]: f32, %[[IN2:.*]]: f32, %[[OUT:.*]]: f32):
109+
// CHECK: %[[VAL_90:.*]] = arith.mulf %[[IN1]], %[[IN2]] : f32
110+
// CHECK: %[[VAL_91:.*]] = arith.addf %[[VAL_90]], %[[OUT]] : f32
111+
/// CHECK: linalg.yield %[[VAL_91]] : f32
110112
// CHECK: } -> tensor<4x4xf32>
111-
// CHECK: %[[S24:.*]] = linalg.mul ins(%[[S23]], %[[S21]] : tensor<4x4xf32>, tensor<4x4xf32>) outs(%[[S22]] : tensor<4x4xf32>) -> tensor<4x4xf32>
112-
// CHECK: %[[S25:.*]] = tensor.extract_slice %[[ARG10]][%[[ARG7]], 0, 0, %[[ARG9]]] [1, 4, 4, 1] [1, 1, 1, 1]
113-
// CHECK: %[[S26:.*]] = tensor.empty() : tensor<4x4xf32>
114-
// CHECK: %[[S27:.*]] = linalg.add ins(%[[S24]], %[[S25]] : tensor<4x4xf32>, tensor<4x4xf32>) outs(%[[S26]] : tensor<4x4xf32>) -> tensor<4x4xf32>
115-
// CHECK: %[[INSERTED_SLICE_9:.*]] = tensor.insert_slice %[[S27]] into %[[ARG10]][%[[ARG7]], 0, 0, %[[ARG9]]] [1, 4, 4, 1] [1, 1, 1, 1]
113+
// CHECK: %[[INSERTED_SLICE_9:.*]] = tensor.insert_slice %[[S23]] into %[[ARG10]][%[[ARG7]], 0, 0, %[[ARG9]]] [1, 4, 4, 1] [1, 1, 1, 1]
116114
// CHECK: scf.yield %[[INSERTED_SLICE_9]]
117115
// CHECK: scf.yield %[[S15]]
118116
// CHECK: %[[S13:.*]] = affine.apply #[[$MAP0]](%[[ARG3]])
@@ -234,22 +232,20 @@ module attributes {transform.with_named_sequence} {
234232
// CHECK: %[[S12:.*]] = scf.for %[[ARG8:.*]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG9:.*]] = %[[EXTRACTED_SLICE_10]])
235233
// CHECK: %[[S15:.*]] = scf.for %[[ARG10:.*]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG11:.*]] = %[[ARG9]])
236234
// CHECK: %[[EXTRACTED_SLICE_11:.*]] = tensor.extract_slice %[[EXTRACTED_SLICE_9]][0, 0, 0, 0, %[[ARG8]], %[[ARG10]]] [6, 6, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1]
235+
// CHECK: %[[S26:.*]] = tensor.extract_slice %[[ARG11]][%[[ARG8]], 0, 0, %[[ARG10]]] [1, 4, 4, 1] [1, 1, 1, 1]
237236
// CHECK: %[[S17:.*]] = tensor.empty() : tensor<4x6xf32>
238237
// CHECK: %[[S18:.*]] = linalg.fill ins(%[[CST_6]] : f32) outs(%[[S17]] : tensor<4x6xf32>) -> tensor<4x6xf32>
239238
// CHECK: %[[S19:.*]] = linalg.matmul ins(%[[CST_1]], %[[EXTRACTED_SLICE_11]] : tensor<4x6xf32>, tensor<6x6xf32>) outs(%[[S18]] : tensor<4x6xf32>) -> tensor<4x6xf32>
240239
// CHECK: %[[S20:.*]] = tensor.empty() : tensor<4x4xf32>
241240
// CHECK: %[[S21:.*]] = linalg.fill ins(%[[CST_6]] : f32) outs(%[[S20]] : tensor<4x4xf32>) -> tensor<4x4xf32>
242241
// CHECK: %[[S22:.*]] = linalg.matmul ins(%[[S19]], %[[CST_0]] : tensor<4x6xf32>, tensor<6x4xf32>) outs(%[[S21]] : tensor<4x4xf32>) -> tensor<4x4xf32>
243-
// CHECK: %[[S23:.*]] = tensor.empty() : tensor<4x4xf32>
244-
// CHECK: %[[S24:.*]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%[[CST]] : f32) outs(%[[S23]] : tensor<4x4xf32>) {
245-
// CHECK: ^bb0(%[[IN:.*]]: f32, %[[OUT:.*]]: f32):
246-
// CHECK: linalg.yield %[[IN]] : f32
242+
// CHECK: %[[S24:.*]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%[[CST]], %[[S22]] : f32, tensor<4x4xf32>) outs(%[[S26]] : tensor<4x4xf32>) {
243+
// CHECK: ^bb0(%[[IN1:.*]]: f32, %[[IN2:.*]]: f32, %[[OUT:.*]]: f32):
244+
// CHECK: %[[VAL_104:.*]] = arith.mulf %[[IN1]], %[[IN2]] : f32
245+
// CHECK: %[[VAL_105:.*]] = arith.addf %[[VAL_104]], %[[OUT]] : f32
246+
/// CHECK: linalg.yield %[[VAL_105]] : f32
247247
// CHECK: } -> tensor<4x4xf32>
248-
// CHECK: %[[S25:.*]] = linalg.mul ins(%[[S24]], %[[S22]] : tensor<4x4xf32>, tensor<4x4xf32>) outs(%[[S23]] : tensor<4x4xf32>) -> tensor<4x4xf32>
249-
// CHECK: %[[S26:.*]] = tensor.extract_slice %[[ARG11]][%[[ARG8]], 0, 0, %[[ARG10]]] [1, 4, 4, 1] [1, 1, 1, 1]
250-
// CHECK: %[[S27:.*]] = tensor.empty() : tensor<4x4xf32>
251-
// CHECK: %[[S28:.*]] = linalg.add ins(%[[S25]], %[[S26]] : tensor<4x4xf32>, tensor<4x4xf32>) outs(%[[S27]] : tensor<4x4xf32>) -> tensor<4x4xf32>
252-
// CHECK: %[[INSERTED_SLICE_12:.*]] = tensor.insert_slice %[[S28]] into %[[ARG11]][%[[ARG8]], 0, 0, %[[ARG10]]] [1, 4, 4, 1] [1, 1, 1, 1]
248+
// CHECK: %[[INSERTED_SLICE_12:.*]] = tensor.insert_slice %[[S24]] into %[[ARG11]][%[[ARG8]], 0, 0, %[[ARG10]]] [1, 4, 4, 1] [1, 1, 1, 1]
253249
// CHECK: scf.yield %[[INSERTED_SLICE_12]]
254250
// CHECK: scf.yield %[[S15]] : tensor<2x4x4x2xf32>
255251
// CHECK: %[[S13:.*]] = affine.apply #[[$MAP0]](%[[ARG4]])
@@ -336,19 +332,17 @@ module attributes {transform.with_named_sequence} {
336332
// CHECK: %[[S6:.*]] = scf.for %[[ARG3:.*]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG4:.*]] = %[[ARG2]])
337333
// CHECK: %[[S7:.*]] = scf.for %[[ARG5:.*]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG6:.*]] = %[[ARG4]])
338334
// CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[EXPANDED]][0, 0, 0, 0, %[[ARG3]], %[[ARG5]]] [6, 1, 1, 1, 1, 1] [1, 1, 1, 1, 1, 1]
335+
// CHECK: %[[S15:.*]] = tensor.extract_slice %[[ARG6]][%[[ARG3]], 0, 0, %[[ARG5]]] [1, 4, 1, 1] [1, 1, 1, 1]
339336
// CHECK: %[[S9:.*]] = tensor.empty() : tensor<4x1xf32>
340337
// CHECK: %[[S10:.*]] = linalg.fill ins(%[[CST_3]] : f32) outs(%[[S9]] : tensor<4x1xf32>) -> tensor<4x1xf32>
341338
// CHECK: %[[S11:.*]] = linalg.matmul ins(%[[CST_0]], %[[EXTRACTED_SLICE]] : tensor<4x6xf32>, tensor<6x1xf32>) outs(%[[S10]] : tensor<4x1xf32>) -> tensor<4x1xf32>
342-
// CHECK: %[[S12:.*]] = tensor.empty() : tensor<4x1xf32>
343-
// CHECK: %[[S13:.*]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel"]} ins(%[[CST]] : f32) outs(%[[S12]] : tensor<4x1xf32>) {
344-
// CHECK: ^bb0(%[[IN:.*]]: f32, %[[OUT:.*]]: f32):
345-
// CHECK: linalg.yield %[[IN]] : f32
339+
// CHECK: %[[S13:.*]] = linalg.generic {indexing_maps = [#map, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%[[CST]], %[[S11]] : f32, tensor<4x1xf32>) outs(%[[S15]] : tensor<4x1xf32>) {
340+
// CHECK: ^bb0(%[[IN1:.*]]: f32, %[[IN2:.*]]: f32, %[[OUT:.*]]: f32):
341+
// CHECK: %[[VAL_57:.*]] = arith.mulf %[[IN1]], %[[IN2]] : f32
342+
// CHECK: %[[VAL_58:.*]] = arith.addf %[[VAL_57]], %[[OUT]] : f32
343+
/// CHECK: linalg.yield %[[VAL_58]] : f32
346344
// CHECK: } -> tensor<4x1xf32>
347-
// CHECK: %[[S14:.*]] = linalg.mul ins(%[[S13]], %[[S11]] : tensor<4x1xf32>, tensor<4x1xf32>) outs(%[[S12]] : tensor<4x1xf32>) -> tensor<4x1xf32>
348-
// CHECK: %[[S15:.*]] = tensor.extract_slice %[[ARG6]][%[[ARG3]], 0, 0, %[[ARG5]]] [1, 4, 1, 1] [1, 1, 1, 1]
349-
// CHECK: %[[S16:.*]] = tensor.empty() : tensor<4x1xf32>
350-
// CHECK: %[[S17:.*]] = linalg.add ins(%[[S14]], %[[S15]] : tensor<4x1xf32>, tensor<4x1xf32>) outs(%[[S16]] : tensor<4x1xf32>) -> tensor<4x1xf32>
351-
// CHECK: %[[INSERTED_SLICE:.*]] = tensor.insert_slice %[[S17]] into %[[ARG6]][%[[ARG3]], 0, 0, %[[ARG5]]] [1, 4, 1, 1] [1, 1, 1, 1]
345+
// CHECK: %[[INSERTED_SLICE:.*]] = tensor.insert_slice %[[S13]] into %[[ARG6]][%[[ARG3]], 0, 0, %[[ARG5]]] [1, 4, 1, 1] [1, 1, 1, 1]
352346
// CHECK: scf.yield %[[INSERTED_SLICE]]
353347
// CHECK: scf.yield %[[S7]]
354348
// CHECK: return %[[S6]]

0 commit comments

Comments
 (0)