[MLIR] Update linalg.conv lowering to use affine load in the absence of padding

[bondhugula]

Summary

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 by affine loads 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 scf/std output obtained after lowering
out affine.

[nicolasvasilache]

Thanks Uday!

[bondhugula]

Minor update comment.

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

let extraClassDeclaration = libraryCallName # [{
     }
 
     // A(i, r_j) * B(r_j) -> C(i)
     llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
       MLIRContext *context = getContext();
       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)
       };
     }
   }];
 
   let hasFolder = 1;
 }
int64_t getDilation(unsigned i) {
         .cast<IntegerAttr>().getValue().getSExtValue();
     }
 
     int64_t getLowPad(unsigned i) {
       assert(i < getNumWindowLoops());
       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});
+    }
   }];
 }
 
 def ConvOp : PoolingBase_Op<"conv", [NInputs<2>, NOutputs<1>]> {
 
   let description = [{
     Generic n-D convolution as described in the TF documentation:
     https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/nn/convolution
let extraClassDeclaration = commonUtils # [{
     // TODO(ntv) extend to support more than 1 dimensions and potentially
     // grouping too.
     unsigned getNumBatchDimensions() { return 1; }
 
     unsigned getNumInputFeatureDimensions() { return 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.
       unsigned nPar = getOutputShapedType(0).getRank();
       unsigned nRed = getNumInputFeatureDimensions();
       // Window loops are a special kind of reduction that is never tiled or
       // parallelized across; i.e. [zs] in the TF notation above whose number
       // match `xs` (i.e. 1 window loop per "image" dimension).

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

static Value getConvOpInput(ConvOp convOp, StdIndexedValue im,
     auto &b = ScopedContext::getBuilderRef();
     Type type = convOp.input().getType().cast<MemRefType>().getElementType();
     Value zero = std_constant(type, b.getZeroAttr(type));
     Value readInput = im(clampedImIdx);
     return conds.empty() ? readInput
                          : (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");
     auto &b = ScopedContext::getBuilderRef();
     auto loc = ScopedContext::getLocation();
     auto mapsRange = convOp.indexing_maps().getAsRange<AffineMapAttr>();
     auto maps = llvm::to_vector<8>(llvm::map_range(
         mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
     SmallVector<Value, 8> fIdx(
         makeCanonicalAffineApplies(b, loc, maps[0], allIvs));
     SmallVector<Value, 8> imIdx(
         makeCanonicalAffineApplies(b, loc, maps[1], allIvs));
     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.
+    // 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);
+    }
   }
 };
 
 template <typename IndexedValueType>
 class LinalgScopedEmitter<IndexedValueType, PoolingMaxOp> {
 public:
   static void emitScalarImplementation(ArrayRef<Value> allIvs,
                                        PoolingMaxOp op) {

mlir/test/Dialect/Linalg/affine.mlir

 //       CHECK:   %[[B:.*]] = dim %arg1, 0 : memref<?x?x?xf32, #[[strided3D]]>
 //       CHECK:   %[[X0:.*]] = dim %arg2, 1 : memref<?x?x?xf32, #[[strided3D]]>
 //       CHECK:   affine.for %{{.*}} = 0 to %[[B]] {
 //       CHECK:     affine.for %{{.*}} = 0 to %[[X0]] {
 //       CHECK:       affine.for %{{.*}} = 0 to %[[K]] {
 //       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>,
                    %arg2: memref<?x?x?x?xf32>) {
   linalg.conv(%arg0, %arg1, %arg2) {dilations = [1, 1],
                                     padding = dense<[[0, 1], [1, 1]]> : tensor<2x2xi64>,
                                     strides = [1, 1]} :
     memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>

 //       CHECK:         affine.for %{{.*}} = 0 to %[[K]] {
 //       CHECK:           affine.for %{{.*}} = 0 to %[[Q]] {
 //       CHECK:             affine.for %{{.*}} = 0 to %[[Z0]] {
 //       CHECK:               affine.for %{{.*}} = 0 to %[[Z1]] {
 //       CHECK:                 %[[SUM0:.*]] = affine.apply #{{.*}}(%{{.*}}, %{{.*}})
 //       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>
 //       CHECK:                 %{{.*}} = mulf %{{.*}}, %{{.*}} : f32
 //       CHECK:                 %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?x?xf32>
 //       CHECK:                 %{{.*}} = addf %{{.*}}, %{{.*}} : f32
 //       CHECK:                 affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?x?xf32>