diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td @@ -269,7 +269,8 @@ let arguments = (ins AnyStridedMemRef:$filter, AnyStridedMemRef:$input, AnyStridedMemRef:$output, OptionalAttr:$strides, - OptionalAttr:$dilations); + OptionalAttr:$dilations, + OptionalAttr:$padding); let extraClassDeclaration = libraryCallName # [{ // TODO(ntv) extend to support more than 1 dimensions and potentially diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp --- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp +++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp @@ -906,8 +906,14 @@ assert(xs.size() == zs.size()); SmallVector res; res.reserve(xs.size()); - for (unsigned i = 0, e = xs.size(); i < e; ++i) - res.push_back(op.getStride(i) * xs[i] + op.getDilation(i) * zs[i]); + for (unsigned i = 0, e = xs.size(); i < e; ++i) { + auto expr = op.getStride(i) * xs[i] + op.getDilation(i) * zs[i]; + if (op.padding()) { + auto padLow = op.padding().getValue().getValue({i, 0}); + expr = expr - padLow.getInt(); + } + res.push_back(expr); + } return res; } diff --git a/mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp b/mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp @@ -175,6 +175,42 @@ template class LinalgScopedEmitter { public: + /// Returns the input value of convOp. If the indices in `imIdx` is out of + /// boundrary, returns 0 instead. + static ValueHandle getConvOpInput(ConvOp convOp, IndexedValueType im, + ArrayRef imIdx) { + if (!convOp.padding()) + return im(imIdx); + + ValueHandle zeroIndex = std_constant_index(0); + SmallVector conds; + for (auto iter : llvm::enumerate(imIdx)) { + using edsc::op::operator<; + using edsc::op::operator>=; + using edsc::op::operator||; + + auto applyOrToLastCond = [&conds](ValueHandle a) { + if (conds.empty()) + conds.push_back(a); + else + conds.push_back(conds.back() || a); + }; + + int idx = iter.index(); + auto dim = iter.value(); + applyOrToLastCond(dim < zeroIndex); + + ValueHandle bound = std_dim(convOp.input(), idx); + applyOrToLastCond(dim >= bound); + } + + auto b = ScopedContext::getBuilder(); + Type type = convOp.input().getType().cast().getElementType(); + ValueHandle zero = std_constant(type, b.getZeroAttr(type)); + ValueHandle readInput = im(imIdx); + return std_select(conds.back(), zero, readInput); + } + static void emitScalarImplementation(ArrayRef allIvs, ConvOp convOp) { assert(convOp.hasBufferSemantics() && "expected linalg op with buffer semantics"); @@ -190,8 +226,10 @@ SmallVector oIdx( makeCanonicalAffineApplies(b, loc, maps[2], allIvs)); IndexedValueType F(convOp.filter()), I(convOp.input()), O(convOp.output()); + // Emit scalar form. - O(oIdx) += F(fIdx) * I(imIdx); + ValueHandle paddedInput = getConvOpInput(convOp, I, imIdx); + O(oIdx) += F(fIdx) * paddedInput; } }; diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp @@ -342,6 +342,9 @@ tileSizes.size() && "expected matching number of tile sizes and loops"); + if (auto convOp = dyn_cast(op.getOperation())) + assert(!convOp.padding() && "expected ConvOp doesn't has padding"); + // If permutation is empty, use the identity. Build the permutation map // otherwise. auto invPermutationMap = AffineMap::getMultiDimIdentityMap( @@ -421,6 +424,9 @@ if (tileSizes.empty()) return llvm::None; + if (auto convOp = dyn_cast(op.getOperation())) + assert(!convOp.padding() && "expected ConvOp doesn't has padding"); + // The following uses the convention that "tiling by zero" skips tiling a // particular dimension. This convention is significantly simpler to handle // instead of adjusting affine maps to account for missing dimensions. diff --git a/mlir/test/Dialect/Linalg/loops.mlir b/mlir/test/Dialect/Linalg/loops.mlir --- a/mlir/test/Dialect/Linalg/loops.mlir +++ b/mlir/test/Dialect/Linalg/loops.mlir @@ -212,6 +212,42 @@ // CHECK: %{{.*}} = addf %{{.*}}, %{{.*}} : f32 // CHECK: store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref +func @conv_padding(%arg0: memref, + %arg1: memref, + %arg2: memref) { + linalg.conv(%arg0, %arg1, %arg2) {dilations = [1, 1], + padding = dense<[[0, 1], [1, 1]]> : tensor<2x2xi64>, + strides = [1, 1]} : + memref, memref, memref + return +} +// CHECK-LABLE: func @conv_padding +// CHECK: %{{.*}}: memref, %{{.*}}: memref, %{{.*}}: memref) { +// CHECK: %[[ZERO:.*]] = constant 0.000000e+00 : f32 +// CHECK: %[[Z0:.*]] = dim %arg0, 0 : memref +// CHECK: %[[Z1:.*]] = dim %arg0, 1 : memref +// CHECK: %[[Q:.*]] = dim %arg0, 2 : memref +// CHECK: %[[K:.*]] = dim %arg0, 3 : memref +// CHECK: %[[B:.*]] = dim %arg1, 0 : memref +// CHECK: %[[X0:.*]] = dim %arg2, 1 : memref +// CHECK: %[[X1:.*]] = dim %arg2, 2 : memref +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[B]] step %{{.*}} { +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[X0]] step %{{.*}} { +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[X1]] step %{{.*}} { +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[K]] step %{{.*}} { +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[Q]] step %{{.*}} { +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[Z0]] step %{{.*}} { +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[Z1]] step %{{.*}} { +// CHECK: %[[SUM0:.*]] = affine.apply #{{.*}}(%{{.*}}, %{{.*}}) +// CHECK: %[[SUM1:.*]] = affine.apply #{{.*}}(%{{.*}}, %{{.*}}) +// CHECK: %{{.*}} = load %{{.*}}[%{{.*}}, %[[SUM0]], %[[SUM1]], %{{.*}}] : memref +// CHECK: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : f32 +// CHECK: %{{.*}} = load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref +// CHECK: %{{.*}} = mulf %{{.*}}, %{{.*}} : f32 +// CHECK: %{{.*}} = load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref +// CHECK: %{{.*}} = addf %{{.*}}, %{{.*}} : f32 +// CHECK: store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref + func @foo(%0: f32, %1: f32, %2: f32) -> (f32, f32) { %f0 = constant 0.0 : f32 return %f0, %f0 : f32, f32 diff --git a/mlir/test/Dialect/Linalg/roundtrip.mlir b/mlir/test/Dialect/Linalg/roundtrip.mlir --- a/mlir/test/Dialect/Linalg/roundtrip.mlir +++ b/mlir/test/Dialect/Linalg/roundtrip.mlir @@ -222,6 +222,28 @@ // ----- +func @conv_padding(%arg0: memref, + %arg1: memref, + %arg2: memref) { + linalg.conv(%arg0, %arg1, %arg2) {dilations = [1, 1], + padding = dense<[[0, 1], [1, 1]]> : tensor<2x2xi64>, + strides = [1, 1]} : + memref, memref, memref + return +} + +// CHECK-LABEL: func @conv_padding( +// CHECK: linalg.conv(%{{.*}}, %{{.*}}, %{{.*}}) { +// CHECK-SAME: dilations = [1, 1], +// CHECK-SAME: padding = dense<[ +// CHECK-SAME: [0, 1], [1, 1]]> : tensor<2x2xi64>, +// CHECK-SAME: strides = [1, 1]} : +// CHECK-SAME: memref, +// CHECK-SAME: memref, +// CHECK-SAME: memref + +// ----- + // CHECK-DAG: #[[strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)> // CHECK-DAG: #[[strided3D:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>