Breadcrumbs where to connect pooling for vectorizations

[nicolasvasilache]

more crumbs

[nicolasvasilache]

served its purpose

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

Value inserted = tensor::createCanonicalRankReducingInsertSliceOp(
       rewriter, loc, conv1DOp.getResult(0), output);
   rewriter.replaceOp(convOp, inserted);
 
   return conv1DOp;
 }
 
 void linalg::populateDecomposeConvolutionPatterns(RewritePatternSet &patterns,
                                                   PatternBenefit benefit) {
+  // TODO: Separately need to add a downscale for pooling.
   patterns.add<DownscaleSizeOneWindowed2DConvolution<linalg::Conv2DNhwcHwcfOp,
                                                      Conv1DNwcWcfOp>,
                DownscaleSizeOneWindowed2DConvolution<linalg::Conv2DNchwFchwOp,
                                                      Conv1DNcwFcwOp>,
                DownscaleDepthwiseConv2DNhwcHwcOp>(patterns.getContext(),
                                                   benefit);
 }

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

Conv1DGenerator(OpBuilder &builder, LinalgOp linalgOp, int strideW,
     // Check for reduction `add` preceded by `mul`.
     Operation *reduceOp = matchLinalgReduction(linalgOp.getDpsInitOperand(0));
     if (!reduceOp)
       return;
     llvm::Optional<vector::CombiningKind> maybeKind;
     maybeKind = getCombinerOpKind(reduceOp);
     if (!maybeKind || *maybeKind != vector::CombiningKind::ADD)
       return;
+
+    // TODO: this is very-specific to conv, and would fail with pooling.
+    // If (region has 2 ops (reduction + yield) and rhs is not used) then it is
+    // the body of a pooling ?
+    // TODO: Rob: there may be a cast i8 -> i32 before the reduction operation.
+
     // Check for single `mul` predecessor. The `mul` operands must be block
     // arguments or extension of block arguments.
     Operation *mulOp = nullptr;
     for (Value operand : reduceOp->getOperands()) {
       if (operand.isa<BlockArgument>())
         continue;
       if (mulOp)
         return;
case Conv1DOpOrder::Nwc:
       bindShapeDims(resShapedType, nSize, wSize);
       lhsShape = {nSize,
                   // iw = ow * sw + kw *  dw - 1
                   //   (i.e. 16 convolved with 3 (@stride 1 dilation 1) -> 14)
                   // Perform the proper inclusive -> exclusive -> inclusive.
                   ((wSize - 1) * strideW + 1) + ((kwSize - 1) * dilationW + 1) -
                       1,
                   cSize};
+      // TODO: irrelevant for pooling
       rhsShape = {kwSize, cSize, fSize};
       resShape = {nSize, wSize, fSize};
       break;
     case Conv1DOpOrder::Ncw:
       // kernel{f, c, kw}
       bindShapeDims(rhsShapedType, fSize, cSize, kwSize);
       // out{n, f, w}
       bindShapeDims(resShapedType, nSize, fSize, wSize);
       lhsShape = {nSize, cSize,
                   // iw = ow * sw + kw *  dw - 1
                   //   (i.e. 16 convolved with 3 (@stride 1 dilation 1) -> 14)
                   // Perform the proper inclusive -> exclusive -> inclusive.
                   ((wSize - 1) * strideW + 1) + ((kwSize - 1) * dilationW + 1) -
                       1};
+      // TODO: irrelevant for pooling
       rhsShape = {fSize, cSize, kwSize};
       resShape = {nSize, fSize, wSize};
       break;
     }
 
     vector::TransferWriteOp write;
     Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
 
     // w is unrolled (i.e. wSizeStep == 1) iff strideW > 1.
     // When strideW == 1, we can batch the contiguous loads and avoid
     // unrolling
     int64_t wSizeStep = strideW == 1 ? wSize : 1;
 
     Type lhsEltType = lhsShapedType.getElementType();
+
+    // TODO: irrelevant for pooling
     Type rhsEltType = rhsShapedType.getElementType();
     Type resEltType = resShapedType.getElementType();
     auto lhsType = VectorType::get(lhsShape, lhsEltType);
     auto rhsType = VectorType::get(rhsShape, rhsEltType);
     auto resType = VectorType::get(resShape, resEltType);
     // Read lhs slice of size {w * strideW + kw * dilationW, c, f} @ [0, 0,
     // 0].
     Value lhs = builder.create<vector::TransferReadOp>(
         loc, lhsType, lhsShaped, ValueRange{zero, zero, zero});
     // Read rhs slice of size {kw, c, f} @ [0, 0, 0].
+
+    // TODO: if pooling do not load
     Value rhs = builder.create<vector::TransferReadOp>(
         loc, rhsType, rhsShaped, ValueRange{zero, zero, zero});
     // Read res slice of size {n, w, f} @ [0, 0, 0].
     Value res = builder.create<vector::TransferReadOp>(
         loc, resType, resShaped, ValueRange{zero, zero, zero});
 
     // The base vectorization case is input: {n,w,c}, weight: {kw,c,f}, output:
     // {n,w,f}. To reuse the base pattern vectorization case, we do pre
     // transpose on input, weight, and output.
     switch (conv1DOpOrder) {
     case Conv1DOpOrder::Nwc:
       // Base case, so no transposes necessary.
       break;
     case Conv1DOpOrder::Ncw: {
       // To match base vectorization case, we pre-transpose current case.
       // ncw -> nwc
       static constexpr std::array<int64_t, 3> permLhs = {0, 2, 1};
       lhs = builder.create<vector::TransposeOp>(loc, lhs, permLhs);
       // fcw -> wcf
       static constexpr std::array<int64_t, 3> permRhs = {2, 1, 0};
+
+      // TODO: irrelevant for pooling
       rhs = builder.create<vector::TransposeOp>(loc, rhs, permRhs);
       // nfw -> nwf
       static constexpr std::array<int64_t, 3> permRes = {0, 2, 1};
       res = builder.create<vector::TransposeOp>(loc, res, permRes);
       break;
     }
     }
 
FailureOr<Operation *> conv(Conv1DOpOrder conv1DOpOrder) {
 
     auto linearIndex = [&](int64_t kw, int64_t w) {
       return kw * (wSize / wSizeStep) + w;
     };
 
     // Compute contraction: O{n, w, f} += I{n, sw * w + dw * kw, c} * F{c, f}
     for (int64_t kw = 0; kw < kwSize; ++kw) {
       for (int64_t w = 0; w < wSize; w += wSizeStep) {
+        // TODO: pool1DSliceAsReduction (maybe contraction)
         resVals[w] = conv1dSliceAsContraction(
             builder, loc, lhsVals[linearIndex(kw, w)], rhsVals[kw], resVals[w]);
       }
     }
 
     // Write back res slice: {n, wSizeStep, f} @ [0, w, 0].
     // This does not depend on kw.
     for (int64_t w = 0; w < wSize; w += wSizeStep) {