[Mlir-commits] [mlir] 0f6999a - [MLIR] Update linalg.conv lowering to use affine load in the absence of padding
Uday Bondhugula
llvmlistbot at llvm.org
Fri Jun 5 00:06:00 PDT 2020
Author: Uday Bondhugula
Date: 2020-06-05T12:28:30+05:30
New Revision: 0f6999af88a08bd430699e72982caa9daf6fa604
URL: https://github.com/llvm/llvm-project/commit/0f6999af88a08bd430699e72982caa9daf6fa604
DIFF: https://github.com/llvm/llvm-project/commit/0f6999af88a08bd430699e72982caa9daf6fa604.diff
LOG: [MLIR] Update linalg.conv lowering to use affine load in the absence of padding
Update linalg to affine lowering for convop to use affine load for input
whenever there is no padding. It had always been using std.loads because
max in index functions (needed for non-zero padding if not materializing
zeros) couldn't be represented in the non-zero padding cases.
In the future, the non-zero padding case could also be made to use
affine - either by materializing or using affine.execute_region. The
latter approach will not impact the scf/std output obtained after
lowering out affine.
Differential Revision: https://reviews.llvm.org/D81191
Added:
Modified:
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
mlir/test/Dialect/Linalg/affine.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index c2784d08b2d2..c13a79087d85 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -215,7 +215,7 @@ def MatvecOp : LinalgStructured_Op<"matvec", [NInputs<2>, NOutputs<1>]> {
AffineExpr i, r_j;
bindDims(context, i, r_j);
return SmallVector<AffineMap, 8>{
- AffineMap::get(2, 0, {i, r_j}, context),
+ AffineMap::get(2, 0, {i, r_j}, context),
AffineMap::get(2, 0, {r_j}, context),
AffineMap::get(2, 0, {i}, context)
};
@@ -314,6 +314,12 @@ class PoolingBase_Op<string mnemonic, list<OpTrait> props>
if (!padding().hasValue()) return 0;
return padding().getValue().getValue<int64_t>({i, 0});
}
+
+ int64_t getHighPad(unsigned i) {
+ assert(i < getNumWindowLoops());
+ if (!padding().hasValue()) return 0;
+ return padding().getValue().getValue<int64_t>({i, 1});
+ }
}];
}
@@ -357,6 +363,11 @@ def ConvOp : PoolingBase_Op<"conv", [NInputs<2>, NOutputs<1>]> {
unsigned getNumOutputFeatureDimensions() { return 1; }
+ unsigned getNumSpatialDimensions() {
+ return getOutputShapedType(0).getRank() - getNumBatchDimensions() -
+ getNumOutputFeatureDimensions();
+ }
+
llvm::Optional<SmallVector<StringRef, 8>> referenceIterators() {
// Outer parallel loops are always the number of output dimensions; i.e.
// [b, xs, q] in the TF notation above.
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp b/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
index 910078875f57..c1080ead1cf9 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
@@ -337,6 +337,15 @@ class LinalgScopedEmitter<IndexedValueType, ConvOp> {
: (Value)std_select(conds.back(), zero, readInput);
}
+ /// Returns true is `convOp` has a non-zero padding.
+ static bool hasPadding(ConvOp convOp) {
+ for (unsigned i = 0, e = convOp.getNumSpatialDimensions(); i < e; ++i) {
+ if (convOp.getLowPad(i) > 0 || convOp.getHighPad(i) > 0)
+ return true;
+ }
+ return false;
+ }
+
static void emitScalarImplementation(ArrayRef<Value> allIvs, ConvOp convOp) {
assert(convOp.hasBufferSemantics() &&
"expected linalg op with buffer semantics");
@@ -352,14 +361,19 @@ class LinalgScopedEmitter<IndexedValueType, ConvOp> {
SmallVector<Value, 8> oIdx(
makeCanonicalAffineApplies(b, loc, maps[2], allIvs));
- // Padded conv involves an affine.max in the memory access which is not
- // allowed by affine.load. Override to always use an StdIndexedValue.
- StdIndexedValue I(convOp.input());
IndexedValueType F(convOp.filter()), O(convOp.output());
- // Emit scalar form.
- Value paddedInput = getConvOpInput(convOp, I, imIdx);
- O(oIdx) += F(fIdx) * paddedInput;
+ // Emit scalar form. Padded conv involves an affine.max in the memory access
+ // which is not allowed by affine.load. Override to use an StdIndexedValue
+ // when there is non-zero padding.
+ if (hasPadding(convOp)) {
+ StdIndexedValue I(convOp.input());
+ Value paddedInput = getConvOpInput(convOp, I, imIdx);
+ O(oIdx) += F(fIdx) * paddedInput;
+ } else {
+ IndexedValueType I(convOp.input());
+ O(oIdx) += F(fIdx) * I(imIdx);
+ }
}
};
diff --git a/mlir/test/Dialect/Linalg/affine.mlir b/mlir/test/Dialect/Linalg/affine.mlir
index e6027c93033a..279a4bde29aa 100644
--- a/mlir/test/Dialect/Linalg/affine.mlir
+++ b/mlir/test/Dialect/Linalg/affine.mlir
@@ -54,6 +54,9 @@ func @conv_view3(%arg0: memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>, %arg1:
// CHECK: affine.for %{{.*}} = 0 to %[[Q]] {
// CHECK: affine.for %{{.*}} = 0 to %[[Z0]] {
// CHECK: %[[SUM:.*]] = affine.apply #[[stride2Dilation1]](%{{.*}}, %{{.*}})
+// No padding needed here; only affine loads.
+// CHECK-NEXT: affine.load
+// CHECK-NEXT: affine.load
func @conv_padding(%arg0: memref<?x?x?x?xf32>,
%arg1: memref<?x?x?x?xf32>,
@@ -85,8 +88,8 @@ func @conv_padding(%arg0: memref<?x?x?x?xf32>,
// CHECK: %[[SUM1:.*]] = affine.apply #{{.*}}(%{{.*}}, %{{.*}})
// CHECK: %[[IDX:.*]] = affine.max #[[clampMinMap]](%[[SUM0]])
// CHECK: %[[IDY:.*]] = affine.max #[[clampMinMap]](%[[SUM1]])
-// Padded conv involves an affine.max in the memory access which is not
-// allowed by affine.load. Override to always use an std.load.
+// Padded conv involves an affine.max in the memory access and this is not
+// allowed by affine.load. Use std.load in such cases.
// CHECK: %{{.*}} = load %{{.*}}[%{{.*}}, %[[IDX]], %[[IDY]], %{{.*}}] : memref<?x?x?x?xf32>
// CHECK: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : f32
// CHECK: %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?x?xf32>
More information about the Mlir-commits
mailing list