[MLIR][Linalg] Hoist padding across multiple levels of tiling

[nicolasvasilache]

Summary

This revision introduces proper backward slice computation during the hoisting of
PadTensorOp. This allows hoisting padding even across multiple levels of tiling.
Such hoisting requires the proper handling of loop bounds that may depend on enclosing
loop variables.

[mehdi_amini]

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

518

debugging leftover

519

Is this something that can be tripped with valid input IR? It isn't clear to me how this invariant is guaranteed by the structure of the callers right now?

571

(ditto)

[nicolasvasilache]

Address review.

[nicolasvasilache]

Address review.

[ftynse]

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

512

s/upper/lower

549–550

This feels like it can lead to invalid IR being produced. Just skipping over an op with regions that produces value will fail to create new values for these ops, and the old values will be used before being defined. I suppose we want to bail out completely. In this case, we also don't want to create ops spuriously, so I would suggest factoring out the validity check (no regions + no side effects) into a separate block/function that runs before we start creating new ops.

552–554

I suppose we shouldn't blindly clone side-effecting ops.

mlir/test/Dialect/Linalg/hoist-padding.mlir

3

Nit: passing -mlir-print-local-scope will print this inline and make tests more self-contained

[nicolasvasilache]

Address review

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

def AffineApplyOp : Affine_Op<"apply", [NoSideEffect]> {
 
   // TODO: The auto-generated builders should check to see if the return type
   // has a constant builder. That way we wouldn't need to explicitly specify the
   // result types here.
   let builders = [
     OpBuilder<(ins "AffineMap":$map, "ValueRange":$mapOperands),
     [{
       build($_builder, $_state, $_builder.getIndexType(), map, mapOperands);
+    }]>,
+    OpBuilder<(ins "ArrayRef<AffineExpr> ":$exprList,"ValueRange":$mapOperands),
+    [{
+      build($_builder, $_state, $_builder.getIndexType(),
+            AffineMap::inferFromExprList(exprList).front(), mapOperands);
     }]>
   ];
 
   let extraClassDeclaration = [{
     /// Returns the affine map to be applied by this operation.
     AffineMap getAffineMap() { return map(); }
 
     /// Returns the affine value map computed from this operation.

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

 //===- Hoisting.cpp - Linalg hoisting transformations ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements functions concerned with hoisting invariant operations
 // in the context of Linalg transformations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
 #include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SCF/Utils.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Dialect/Vector/VectorUtils.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
func.walk([&](vector::TransferReadOp transferRead) {
       loop.erase();
       // Need to interrupt and restart because erasing the loop messes up the
       // walk.
       return WalkResult::interrupt();
     });
   }
 }
 
+/// Return success if `v` is a value that is only transitively defined by ops of
+/// type in `OpTypeList`.
+template <typename... OpTypeList>
+static bool backwardsSliceOnlyHasOpsOfType(scf::ForOp outerLimit, Value v) {
+  // Compute a backward slice up to, but not including, `outerLimit`.
+  llvm::SetVector<Operation *> backwardSlice;
+  getBackwardSlice(v, &backwardSlice, [&](Operation *op) {
+    return outerLimit->isProperAncestor(op);
+  });
+  // Traverse the backward slice and ensure we can perform the computation to
+  // hoist.
+  for (Operation *op : backwardSlice) {
+    if (isa<OpTypeList...>(op))
+      continue;
+    LLVM_DEBUG(DBGS() << "Abort: unadmissible op in slice " << *op << "\n");
+    return false;
+  }
+  return true;
+}
+
+bool isDefinedOutsideOrConstant(scf::ForOp outer, Value v) {
+  return outer.isDefinedOutsideOfLoop(v) || v.getDefiningOp<ConstantOp>();
+}
+
+/// Compute the tightest lower bound with quantities that are all defined
+/// outside of `outer`.
+/// Return null if such a bound cannot be computed.
+Value computeLoopIndependentLowerBound(OpBuilder &b, scf::ForOp outer,
+                                       Value v) {
+  if (isDefinedOutsideOrConstant(outer, v))
+    return v;
+  return Value();
+}
+
+/// Compute the tightest upper bound with quantities that are all defined
+/// outside of `outer`.
+/// Expects all ops in the backward slice of `v` up to `outer` to be either
+/// scf.for, affine.min or affine.apply.
+static Value computeLoopIndependentUpperBound(OpBuilder &b, scf::ForOp outer,
+                                              Value v) {
+  if (isDefinedOutsideOrConstant(outer, v))
+    return v;
+
+  LLVM_DEBUG(DBGS() << "Begin loopIndependentUpperBound for: " << v << "\n");
+
+  bool ok =
+      backwardsSliceOnlyHasOpsOfType<scf::ForOp, AffineMinOp, AffineApplyOp>(
+          outer, v);
+  assert(ok && "expected to only be defined by scf::ForOp and AffineMinOp");
+
+  // Compute a backward slice up to, but not including, `outer`.
+  llvm::SetVector<Operation *> backwardSlice;
+  getBackwardSlice(v, &backwardSlice,
+                   [&](Operation *op) { return outer->isProperAncestor(op); });
+  backwardSlice.insert(v.getDefiningOp());
+
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPoint(outer);
+  Value res = v;
+  BlockAndValueMapping bvm;
+  for (Operation *op : backwardSlice) {
+    if (isa<scf::ForOp>(op))
+      continue;
+    if (isa<AffineApplyOp>(op)) {
+      b.clone(*op, bvm);
+      continue;
+    }
+    auto sliceMinOp = cast<AffineMinOp>(op);
+    // Perform the substitution of the operands of AffineMinOp.
+    auto mapAndOperands = substituteMin(
+        sliceMinOp, [&](Operation *op) { return outer->isAncestor(op); });
+    SmallVector<Value> resultOperands = mapAndOperands.dims;
+    llvm::append_range(resultOperands, mapAndOperands.symbols);
+    AffineMap map = mapAndOperands.map;
+    canonicalizeMapAndOperands(&map, &resultOperands);
+    map = simplifyAffineMap(map);
+    res = b.create<AffineMinOp>(
+        outer->getLoc(), map,
+        llvm::to_vector<4>(llvm::map_range(resultOperands, [&](Value operand) {
+          return bvm.lookupOrDefault(operand);
+        })));
+    bvm.map(sliceMinOp, res);
+  }
+  LLVM_DEBUG(DBGS() << "End loopIndependentUpperBound with: " << res << "\n");
+  return res;
+}
+
+/// Return the number of iterations in the loop (ub - lb).ceilDiv(step).
+/// The returned Value is guaranteed not to depend on any loop comprised in
+/// [`outer`, `forOp`].
+/// Return null if such a loop-independent quantity cannot be computed.
+static Value buildLoopTripCount(OpBuilder &b, scf::ForOp outer,
+                                scf::ForOp forOp) {
+  MLIRContext *ctx = forOp->getContext();
+  AffineExpr lb, ub, step;
+  bindDims(ctx, lb, ub);
+  bindSymbols(ctx, step);
+  Value lbVal = computeLoopIndependentLowerBound(b, outer, forOp.lowerBound()),
+        ubVal = computeLoopIndependentUpperBound(b, outer, forOp.upperBound()),
+        stepVal = forOp.step();
+  if (!lbVal || !ubVal || !stepVal)
+    return Value();
+  auto loc = forOp->getLoc();
+  Value res = b.create<AffineApplyOp>(loc, (ub - lb).ceilDiv(step),
+                                      ValueRange{lbVal, ubVal, stepVal});
+  return res;
+}
+
+/// Return the current iteration number in the loop (iv - lb).ceilDiv(step).
+/// The returned Value is guaranteed not to depend on any loop comprised in
+/// [`outer`, `forOp`].
+/// Return null if such a loop-independent quantity cannot be computed.
+static Value buildLoopIterationCount(OpBuilder &b, scf::ForOp outer,
+                                     scf::ForOp forOp) {
+  MLIRContext *ctx = forOp->getContext();
+  AffineExpr iv, lb, step;
+  bindDims(ctx, iv, lb);
+  bindSymbols(ctx, step);
+  Value ivVal = forOp.getInductionVar(),
+        lbVal = computeLoopIndependentLowerBound(b, outer, forOp.lowerBound()),
+        stepVal = forOp.step();
+  if (!ivVal || !lbVal || !stepVal)
+    return Value();
+  auto loc = forOp->getLoc();
+  return b.create<AffineApplyOp>(loc, (iv - lb).ceilDiv(step),
+                                 ValueRange{ivVal, lbVal, stepVal});
+}
+
 /// Ensure prerequisites that guarantee pad op hoisting can occur.
 /// Return failure in the cases when we cannot perform hoisting; i.e. if either:
 ///   1. There exists a use of `padTensorOp` that is not a linalg input operand.
 ///   2. There isn't an enclosing `outermostEnclosingForOp` loop.
 ///   3. There exists an op with a region that is dominated by
 ///   `outermostEnclosingForOp` and that isn't a LoopLikeInterface or a
 ///    LinalgOp.
-///   3. There exists an op with side effects that is dominated by
-///    `outermostEnclosingForOp` and that isn't a LoopLikeInterface.
+///   4. There exists an op with side effects that is dominated by
+///   `outermostEnclosingForOp` and that isn't a LoopLikeInterface.
+///   5. The lower bound, upper bound and step of all the loops involved in the
+///   hoisting can be
 ///
 /// While ensuring prerequisites:
 ///   1. Fill the `backwardSlice` to contain the topologically sorted ops
 ///   dominated by `outermostEnclosingForOp`.
 ///   2. Fill the `packingLoops` to contain only the enclosing loops of
 ///   `backwardSlice` whose IV is actually used in computing padding. Loops that
 ///   remain in `backwardSlice` but that are not in `packingLoops` are
 ///   dimensions of reuse.
 static LogicalResult
 hoistPaddingOnTensorsPrerequisites(linalg::PadTensorOp padTensorOp, int nLevels,
                                    llvm::SetVector<Operation *> &backwardSlice,
-                                   llvm::SetVector<Operation *> &packingLoops) {
+                                   llvm::SetVector<Operation *> &packingLoops,
+                                   SmallVector<Value> &dynamicTensorSizes) {
   // Bail on any use that isn't an input of a Linalg op.
   // Hoisting of inplace updates happens after vectorization.
   for (OpOperand &use : padTensorOp.result().getUses()) {
     auto linalgUser = dyn_cast<linalg::LinalgOp>(use.getOwner());
     if (!linalgUser || !linalgUser.isInputTensor(&use))
       return failure();
   }
 
for (Operation *user : forOp.getInductionVar().getUsers()) {
       }
     }
   }
 
   // Backward slice is a topologically sorted list of ops starting at
   // `outermostEnclosingForOp`.
   assert(outermostEnclosingForOp == backwardSlice.front());
 
-  return success();
-}
-
-/// Return the number of iterations in the loop (ub - lb).ceilDiv(step).
-static Value buildLoopTripCount(OpBuilder &b, scf::ForOp forOp) {
-  MLIRContext *ctx = forOp->getContext();
-  AffineExpr lb, ub, step;
-  bindDims(ctx, lb, ub);
-  bindSymbols(ctx, step);
-  return b.create<AffineApplyOp>(
-      forOp->getLoc(), AffineMap::get(2, 1, {(ub - lb).ceilDiv(step)}, ctx),
-      ValueRange{forOp.lowerBound(), forOp.upperBound(), forOp.step()});
-}
+  scf::ForOp outer = cast<scf::ForOp>(outermostEnclosingForOp);
+  if (llvm::any_of(packingLoops, [&](Operation *op) {
+        scf::ForOp forOp = cast<scf::ForOp>(op);
+        Value lb = forOp.lowerBound(), ub = forOp.upperBound(),
+              step = forOp.step();
+        return !isDefinedOutsideOrConstant(outer, lb) ||
+               !(isDefinedOutsideOrConstant(outer, ub) ||
+                 backwardsSliceOnlyHasOpsOfType<scf::ForOp, AffineMinOp,
+                                                AffineApplyOp>(outer, ub)) ||
+               !isDefinedOutsideOrConstant(outer, step);
+      }))
+    return failure();
 
-/// Return the current iteration number in the loop (iv - lb).ceilDiv(step).
-static Value buildLoopIterationCount(OpBuilder &b, scf::ForOp forOp) {
-  MLIRContext *ctx = forOp->getContext();
-  AffineExpr iv, lb, step;
-  bindDims(ctx, iv, lb);
-  bindSymbols(ctx, step);
-  return b.create<AffineApplyOp>(
-      forOp->getLoc(), AffineMap::get(2, 1, {(iv - lb).ceilDiv(step)}, ctx),
-      ValueRange{forOp.getInductionVar(), forOp.lowerBound(), forOp.step()});
+  // IP just before the outermost loop considered that we hoist above.
+  OpBuilder b(outermostEnclosingForOp);
+  dynamicTensorSizes =
+      llvm::to_vector<4>(llvm::map_range(packingLoops, [&](Operation *op) {
+        return buildLoopTripCount(b, cast<scf::ForOp>(outermostEnclosingForOp),
+                                  cast<scf::ForOp>(op));
+      }));
+  // Assert all loop trip counts can be computed.
+  if (!llvm::all_of(dynamicTensorSizes, [](Value v) { return v; }))
+    llvm_unreachable("loop independence prerequisite not met");
+  return success();
 }
 
 LogicalResult mlir::linalg::hoistPaddingOnTensors(PadTensorOp &padTensorOp,
                                                   unsigned nLoops) {
+  SmallVector<Value> dynamicTensorSizes;
   llvm::SetVector<Operation *> backwardSlice, packingLoops;
   if (failed(hoistPaddingOnTensorsPrerequisites(padTensorOp, nLoops,
-                                                backwardSlice, packingLoops)))
+                                                backwardSlice, packingLoops,
+                                                dynamicTensorSizes)))
     return failure();
 
   // Update actual number of loops, which may be smaller.
   nLoops = packingLoops.size();
 
   Location loc = padTensorOp->getLoc();
   RankedTensorType paddedTensorType = padTensorOp.getResultType();
   unsigned paddedRank = paddedTensorType.getRank();
 
   // Backward slice is a topologically sorted list of ops starting at
   // `outermostEnclosingForOp`.
   Operation *outermostEnclosingForOp = backwardSlice.front();
   // IP just before the outermost loop considered that we hoist above.
   OpBuilder b(outermostEnclosingForOp);
 
   // Create the packed tensor<?x?x..?xpadded_shape> into which we amortize
   // padding.
   SmallVector<int64_t> packedShape(nLoops, ShapedType::kDynamicSize);
   // TODO: go grab dims when necessary, for now PadTensorOp returns a static
   // tensor.
   llvm::append_range(packedShape, paddedTensorType.getShape());
   auto packedTensorType =
       RankedTensorType::get(packedShape, paddedTensorType.getElementType());
-  auto dynamicSizes =
-      llvm::to_vector<4>(llvm::map_range(packingLoops, [&](Operation *op) {
-        return buildLoopTripCount(b, cast<scf::ForOp>(op));
-      }));
   Value packedTensor = b.create<linalg::InitTensorOp>(
-      loc, dynamicSizes, packedTensorType.getShape(),
+      loc, dynamicTensorSizes, packedTensorType.getShape(),
       packedTensorType.getElementType());
 
   // Clone the operations involved in the backward slice, iteratively stepping
   // into the loops that we encounter.
   // The implementation proceeds in a stack-like fashion:
   //   1. Iteratively clone and step into the loops, pushing the `packedTensor`
   //      deeper in the stack.
   //   2. Create a SubTensorInsert at the top of the stack.
   //   3. Iteratively pop and yield the result of the SubTensorInsertOp across
   //     the cloned loops.
   SmallVector<Value> clonedLoopIvs, leadingPackedTensorIndexings;
   clonedLoopIvs.reserve(nLoops);
   leadingPackedTensorIndexings.reserve(nLoops);
   BlockAndValueMapping bvm;
-  // Stack step 1. iteratively clone loops and push `packedTensor`.
   // Insert `padTensorOp` into the backwardSlice so we clone it too.
   backwardSlice.insert(padTensorOp);
+  // Stack step 1. iteratively clone loops and push `packedTensor`.
   for (Operation *op : backwardSlice) {
     if (op->getNumRegions() == 0 || isa<linalg::PadTensorOp>(op)) {
       b.clone(*op, bvm);
       continue;
     }
     // TODO: support more cases as they appear.
     auto forOp = dyn_cast<scf::ForOp>(op);
     assert(forOp && "Expected scf::ForOp when hoisting pad ops");
     // Unused loop, just skip it.
     if (!packingLoops.contains(forOp))
       continue;
+
     auto clonedForOp =
-        b.create<scf::ForOp>(loc, forOp.lowerBound(), forOp.upperBound(),
-                             forOp.step(), packedTensor);
+        b.create<scf::ForOp>(loc, bvm.lookupOrDefault(forOp.lowerBound()),
+                             bvm.lookupOrDefault(forOp.upperBound()),
+                             bvm.lookupOrDefault(forOp.step()), packedTensor);
+
+    bvm.map(forOp.getInductionVar(), clonedForOp.getInductionVar());
     assert(clonedForOp->getNumRegions() == 1);
     clonedLoopIvs.push_back(clonedForOp.getInductionVar());
+
     b.setInsertionPointToStart(&clonedForOp->getRegion(0).front());
-    leadingPackedTensorIndexings.push_back(
-        buildLoopIterationCount(b, clonedForOp));
-    bvm.map(forOp.getInductionVar(), clonedLoopIvs.back());
+    Value loopIndependentIterationCount = buildLoopIterationCount(
+        b, cast<scf::ForOp>(outermostEnclosingForOp), clonedForOp);
+    // Assert the loop-independent iteration count can be computed.
+    if (!loopIndependentIterationCount)
+      llvm_unreachable("loop independence prerequisite not met");
+    leadingPackedTensorIndexings.push_back(loopIndependentIterationCount);
     packedTensor = clonedForOp.getRegionIterArgs().front();
   }
 
   // Stack step 2. create SubTensorInsertOp at the top of the stack.
   // offsets = [clonedLoopIvs, 0 .. 0].
   SmallVector<OpFoldResult> offsets(leadingPackedTensorIndexings.begin(),
                                     leadingPackedTensorIndexings.end());
   offsets.append(paddedRank, b.getIndexAttr(0));
for (Value iv : llvm::reverse(clonedLoopIvs)) {
     valueToYield = forOp.getResult(0);
   }
 
   // Now the packed tensor is ready, replace the original padding op by a
   // 1x..x1 SubTensor [originalLoopIvs, 0 .. 0][1 .. 1, paddedShape][1 .. 1].
   b.setInsertionPoint(padTensorOp);
   SmallVector<Value> loopIterationCounts =
       llvm::to_vector<4>(llvm::map_range(packingLoops, [&](Operation *loop) {
-        return buildLoopIterationCount(b, cast<scf::ForOp>(loop));
+        return buildLoopIterationCount(
+            b, cast<scf::ForOp>(outermostEnclosingForOp),
+            cast<scf::ForOp>(loop));
       }));
+  // Assert all loop iteration counts can be computed.
+  if (llvm::any_of(loopIterationCounts, [](Value v) { return !v; }))
+    llvm_unreachable("loop independence prerequisite not met");
   // offsets = [originalLoopIvs, 0 .. 0].
   offsets.assign(loopIterationCounts.begin(), loopIterationCounts.end());
   offsets.append(paddedRank, b.getIndexAttr(0));
   // sizes = [1 .. 1, paddedShape] (definedabove).
   // strides = [1 .. 1] (defined above)
   packedTensor =
       scf::getForInductionVarOwner(clonedLoopIvs.front())->getResult(0);
   padTensorOp.replaceAllUsesWith(

mlir/test/Dialect/Linalg/hoist-padding.mlir

-// RUN: mlir-opt %s -test-linalg-transform-patterns=test-hoist-padding-2-level -canonicalize | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -test-linalg-transform-patterns=test-hoist-padding-2-level -canonicalize | FileCheck %s
 
+// CHECK-DAG: #[[$DIV3:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 3)>
+// CHECK-DAG: #[[$DIV4:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 4)>
+// CHECK-DAG: #[[$DIVS3:[0-9a-z]+]] = affine_map<()[s0] -> (s0 ceildiv 3)>
+// CHECK-DAG: #[[$DIVS4:[0-9a-z]+]] = affine_map<()[s0] -> (s0 ceildiv 4)>
 #map0 = affine_map<(d0)[s0] -> (2, -d0 + s0)>
 #map1 = affine_map<(d0)[s0] -> (4, -d0 + s0)>
 #map2 = affine_map<(d0)[s0] -> (3, -d0 + s0)>
 #map3 = affine_map<(d0, d1) -> (2, d0 - d1)>
 #map4 = affine_map<(d0, d1) -> (3, d0 - d1)>
 
-// CHECK-DAG: #[[$DIV3:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 3)>
-// CHECK-DAG: #[[$DIV4:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 4)>
-// CHECK-DAG: #[[$DIVS3:[0-9a-z]+]] = affine_map<()[s0] -> (s0 ceildiv 3)>
-// CHECK-DAG: #[[$DIVS4:[0-9a-z]+]] = affine_map<()[s0] -> (s0 ceildiv 4)>
-
 // CHECK-LABEL: func @matmul_tensors
 //  CHECK-SAME:   %[[TA:[0-9a-z]+]]: tensor
 //  CHECK-SAME:   %[[TB:[0-9a-z]+]]: tensor
 //  CHECK-SAME:   %[[TC:[0-9a-z]+]]: tensor
 func @matmul_tensors(
   %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
   -> tensor<?x?xf32>
 {
%4 = scf.for %arg5 = %c0 to %2 step %c3 iter_args(%arg6 = %arg4) -> (tensor<?x?xf32>) {
         scf.yield %32 : tensor<?x?xf32>
       }
       scf.yield %5 : tensor<?x?xf32>
     }
     scf.yield %4 : tensor<?x?xf32>
   }
   return %3 : tensor<?x?xf32>
 }
+
+// -----
+
+// CHECK-DAG: #[[$MIN_REST8:[0-9a-z]+]] = affine_map<(d0)[s0] -> (8, -d0 + s0)>
+// CHECK-DAG: #[[$MIN_MOD4:[0-9a-z]+]] = affine_map<(d0) -> (4, d0 - ((d0 - 1) floordiv 4) * 4)>
+// CHECK-DAG: #[[$DIV4:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 4)>
+// CHECK-DAG: #[[$DIV2:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 2)>
+#map0 = affine_map<(d0)[s0] -> (8, -d0 + s0)>
+#map1 = affine_map<(d0, d1) -> (4, d0 - d1)>
+#map2 = affine_map<(d0, d1) -> (2, d0 - d1)>
+
+// CHECK-LABEL: func @dot
+func @dot(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>, %arg2: tensor<f32>)
+    -> tensor<f32>
+{
+  %c8 = constant 8 : index
+  %c4 = constant 4 : index
+  %cst = constant 0.000000e+00 : f32
+  %c2 = constant 2 : index
+  %c0 = constant 0 : index
+  %1 = memref.dim %arg0, %c0 : tensor<?xf32>
+  %2 = memref.dim %arg0, %c0 : tensor<?xf32>
+  %3 = memref.dim %arg1, %c0 : tensor<?xf32>
+
+  //      CHECK: scf.for %[[I:[0-9a-z]+]] =
+  //
+  //      CHECK:   %[[MR8:.*]] = affine.min #[[$MIN_REST8]](%[[I]])
+  //      CHECK:   %[[D0:.*]] = affine.apply #[[$DIV4]](%[[MR8]])
+  //      CHECK:   %[[MM4:.*]] = affine.min #[[$MIN_MOD4]](%[[MR8]])
+  //      CHECK:   %[[D1:.*]] = affine.apply #[[$DIV2]](%[[MM4]])
+  // Init tensor and pack.
+  //      CHECK:   %[[INIT_PACKED_A:.*]] = linalg.init_tensor [%[[D0]], %[[D1]], 2] : tensor<?x?x2xf32>
+  //      CHECK:   %[[PACKED_A:.*]] = scf.for %[[II:[0-9a-z]+]] = {{.*}} iter_args(%{{.*}} = %[[INIT_PACKED_A]]) -> (tensor<?x?x2xf32>) {
+  //      CHECK:     scf.for %[[III:[0-9a-z]+]] =
+  //      CHECK:       subtensor_insert %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}, 0] [1, 1, 2] [1, 1, 1] : tensor<2xf32> into tensor<?x?x2xf32>
+  //
+  //      CHECK:   %[[D0_2:.*]] = affine.apply #[[$DIV4]](%[[MR8]])
+  //      CHECK:   %[[MM4_2:.*]] = affine.min #[[$MIN_MOD4]](%[[MR8]])
+  //      CHECK:   %[[D1_2:.*]] = affine.apply #[[$DIV2]](%[[MM4_2]])
+  // Init tensor and pack.
+  //      CHECK:   %[[INIT_PACKED_B:.*]] = linalg.init_tensor [%[[D0_2]], %[[D1_2]], 2] : tensor<?x?x2xf32>
+  //      CHECK:   %[[PACKED_B:.*]] = scf.for %[[II_2:[0-9a-z]+]] = {{.*}} iter_args(%{{.*}} = %[[INIT_PACKED_B]]) -> (tensor<?x?x2xf32>) {
+  //      CHECK:     scf.for %[[III_2:[0-9a-z]+]] =
+  //      CHECK:       subtensor_insert %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}, 0] [1, 1, 2] [1, 1, 1] : tensor<2xf32> into tensor<?x?x2xf32>
+  // Compute.
+  //      CHECK:   scf.for %[[II_3:[0-9a-z]+]] =
+  //      CHECK:     scf.for %[[III_3:[0-9a-z]+]] = {{.*}} iter_args(%[[C:.*]] = %{{.*}}) -> (tensor<f32>) {
+  //      CHECK:       %[[IDX0:.*]] = affine.apply #[[$DIV4]](%[[II_3]])
+  //      CHECK:       %[[IDX1:.*]] = affine.apply #[[$DIV2]](%[[III_3]])
+  //      CHECK:       %[[A:.*]] = subtensor %[[PACKED_A]][%[[IDX0]], %[[IDX1]], 0] [1, 1, 2] [1, 1, 1] : tensor<?x?x2xf32> to tensor<2xf32>
+  //      CHECK:       %[[IDX0_2:.*]] = affine.apply #[[$DIV4]](%[[II_3]])
+  //      CHECK:       %[[IDX1_2:.*]] = affine.apply #[[$DIV2]](%[[III_3]])
+  //      CHECK:       %[[B:.*]] = subtensor %[[PACKED_B]][%[[IDX0_2]], %[[IDX1_2]], 0] [1, 1, 2] [1, 1, 1] : tensor<?x?x2xf32> to tensor<2xf32>
+  //      CHECK:       linalg.dot ins(%[[A]], %[[B]] : tensor<2xf32>, tensor<2xf32>) outs(%[[C]] : tensor<f32>) -> tensor<f32>
+
+  %4 = scf.for %arg3 = %c0 to %1 step %c8 iter_args(%arg4 = %arg2) -> (tensor<f32>) {
+    %5 = affine.min #map0(%arg3)[%2]
+    %6 = subtensor %arg0[%arg3] [%5] [1] : tensor<?xf32> to tensor<?xf32>
+    %7 = affine.min #map0(%arg3)[%3]
+    %8 = subtensor %arg1[%arg3] [%7] [1] : tensor<?xf32> to tensor<?xf32>
+    %9 = scf.for %arg5 = %c0 to %5 step %c4 iter_args(%arg6 = %arg4) -> (tensor<f32>) {
+      %10 = affine.min #map1(%5, %arg5)
+      %11 = subtensor %6[%arg5] [%10] [1] : tensor<?xf32> to tensor<?xf32>
+      %12 = affine.min #map1(%7, %arg5)
+      %13 = subtensor %8[%arg5] [%12] [1] : tensor<?xf32> to tensor<?xf32>
+      %14 = scf.for %arg7 = %c0 to %10 step %c2 iter_args(%arg8 = %arg6) -> (tensor<f32>) {
+        %15 = affine.min #map2(%10, %arg7)
+        %16 = subtensor %11[%arg7] [%15] [1] : tensor<?xf32> to tensor<?xf32>
+        %17 = affine.min #map2(%12, %arg7)
+        %18 = subtensor %13[%arg7] [%17] [1] : tensor<?xf32> to tensor<?xf32>
+        %19 = subi %c2, %15 : index
+        %20 = linalg.pad_tensor %16 low[%c0] high[%19]  {
+        ^bb0(%arg9: index):  // no predecessors
+          linalg.yield %cst : f32
+        } : tensor<?xf32> to tensor<2xf32>
+        %21 = subi %c2, %17 : index
+        %22 = linalg.pad_tensor %18 low[%c0] high[%21]  {
+        ^bb0(%arg9: index):  // no predecessors
+          linalg.yield %cst : f32
+        } : tensor<?xf32> to tensor<2xf32>
+        %23 = linalg.dot ins(%20, %22 : tensor<2xf32>, tensor<2xf32>) outs(%arg8 : tensor<f32>) -> tensor<f32>
+        scf.yield %23 : tensor<f32>
+      }
+      scf.yield %14 : tensor<f32>
+    }
+    scf.yield %9 : tensor<f32>
+  }
+  return %4 : tensor<f32>
+}