[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

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index c2784d0..c13a790 100644 (file)
--- 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 9100788..c1080ea 100644 (file)
--- a/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
@@ -337,6 +337,15 @@ public:
                          : (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 @@ public:
     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 e6027c9..279a4bd 100644 (file)
--- 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>