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,45 @@ template class LinalgScopedEmitter { public: + static ValueHandle getConvOpInput(ConvOp convOp, IndexedValueType I, + ArrayRef imIdx) { + if (!convOp.padding()) + return I(imIdx); + + SmallVector shapes; + Value inputArg = convOp.input(); + auto inputArgShape = inputArg.getType().cast().getShape(); + shapes.append(inputArgShape.begin(), inputArgShape.end()); + + 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); + + assert(shapes[idx] != -1 && "dynamic shape is not available"); + ValueHandle shape = std_constant_index(shapes[idx]); + applyOrToLastCond(dim >= shape); + } + auto b = ScopedContext::getBuilder(); + Type type = convOp.input().getType().cast().getElementType(); + ValueHandle zero = std_constant(type, b.getZeroAttr(type)); + ValueHandle readInput = I(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 +229,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,41 @@ // CHECK: %{{.*}} = addf %{{.*}}, %{{.*}} : f32 // CHECK: store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref +func @conv_padding(%arg0: memref<2x3x1x1xf32>, + %arg1: memref<1x4x5x1xf32>, + %arg2: memref<1x4x5x1xf32>) { + linalg.conv(%arg0, %arg1, %arg2) {dilations = [1, 1], + padding = dense<[[0, 1], [1, 1]]> : tensor<2x2xi64>, + strides = [1, 1]} : + memref<2x3x1x1xf32>, memref<1x4x5x1xf32>, memref<1x4x5x1xf32> + return +} +// CHECK-LABLE: func @conv_padding +// CHECK: %{{.*}}: memref<2x3x1x1xf32>, %{{.*}}: memref<1x4x5x1xf32>, %{{.*}}: memref<1x4x5x1xf32>) { +// CHECK: %[[C2:.*]] = constant 2 : index +// CHECK: %[[C3:.*]] = constant 3 : index +// CHECK: %[[C0:.*]] = constant 0 : index +// CHECK: %[[C4:.*]] = constant 4 : index +// CHECK: %[[C5:.*]] = constant 5 : index +// CHECK: %[[C1:.*]] = constant 1 : index +// CHECK: %[[ZERO:.*]] = constant 0.000000e+00 : f32 +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[C1]] step %{{.*}} { +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[C4]] step %{{.*}} { +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[C5]] step %{{.*}} { +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[C1]] step %{{.*}} { +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[C1]] step %{{.*}} { +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[C2]] step %{{.*}} { +// CHECK: loop.for %{{.*}} = %{{.*}} to %[[C3]] step %{{.*}} { +// CHECK: %[[SUM0:.*]] = affine.apply #{{.*}}(%{{.*}}, %{{.*}}) +// CHECK: %[[SUM1:.*]] = affine.apply #{{.*}}(%{{.*}}, %{{.*}}) +// CHECK: %{{.*}} = load %{{.*}}[%{{.*}}, %[[SUM0]], %[[SUM1]], %{{.*}}] : memref<1x4x5x1xf32> +// CHECK: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : f32 +// CHECK: %{{.*}} = load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<2x3x1x1xf32> +// CHECK: %{{.*}} = mulf %{{.*}}, %{{.*}} : f32 +// CHECK: %{{.*}} = load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<1x4x5x1xf32> +// CHECK: %{{.*}} = addf %{{.*}}, %{{.*}} : f32 +// CHECK: store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<1x4x5x1xf32> + 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<2x3x1x1xf32>, + %arg1: memref<1x4x5x1xf32>, + %arg2: memref<1x4x5x1xf32>) { + linalg.conv(%arg0, %arg1, %arg2) {dilations = [1, 1], + padding = dense<[[0, 1], [1, 1]]> : tensor<2x2xi64>, + strides = [1, 1]} : + memref<2x3x1x1xf32>, memref<1x4x5x1xf32>, memref<1x4x5x1xf32> + 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<2x3x1x1xf32>, +// CHECK-SAME: memref<1x4x5x1xf32>, +// CHECK-SAME: memref<1x4x5x1xf32> + +// ----- + // 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)>