[mlir] Loop bounds inference in linalg.generic op improved to support bounds for convolution

[limo1996]

Summary

Loop bound inference is right now very limited as it supports only permutation maps and thus
it is impossible to implement convolution with linalg.generic as it requires more advanced
loop bound inference. This commits solves it for the convolution case.

Depends On D83158

[ftynse]

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

79

Nit: Top-level functions should use /// for comments.

83

Please explain the semantics of the argument in the top-level comment instead.

84

Prefer ValueRange unless it is missing some functionality

85

We tend to use assert(condition && "textual description")

95

Prefer early-return to decrease horizontal alignment. E.g.,

for (...) {
  if (auto x = dyn_cast<..>) {
    // code
  }
}

can be transformed into

for (...) {

auto x = dyn_cast<..>;
if (!x) continue;
// code

}

97

cast assume the operand has the expected type and fails an assertion if it does not. If there is no other check, e.g. in the op verifier, this may trigger an assertion on user-entered IR, which we must avoid

99

MLIR uses camelBack for identifiers

538

Please use a more descriptive variable name. Also avoid C-style casts.

541

No need for a cast here, but I'd check or assert somewhere (depending on whether it can happen with user-entered IR or not) that numIn > numDims, one can have a map of the shape (i) -> (i,i,i).

[limo1996]

All comments of ftynse incorporated

[ftynse]

Deferring final approval to @nicolasvasilache

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

108

This is an anti-pattern in LLVM, if you intend to cast, you should use dyn_cast and check that the result is non-null instead.

auto llhs = lhs.getLHS().dyn_cast<AffineDimExpr>();
if (!llhs)
  continue;

int dimPosition = llhs.getPosition();

(This will not work for "add" and "mul" above because they don't have a dedicated type)

Also, avoid single-character variable names.

[limo1996]

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

541

numInputs = numDims + numSymbols so it should never happen

[limo1996]

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

551

@ftynse inversePermutation needs to "support symbols by ignoring them" as verification uses the function to check whether concatination of indexing maps is invertible.. Another option is to ignore the check if map has symbols which can cause problems..

[limo1996]

Proposed change to circumvent call to inversePermutation during verification of linalg.generic

[limo1996]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

298 ↗

(On Diff #275972)

@ftynse Here is the proposed change which requires the least amount of code..

[ftynse]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

298 ↗

(On Diff #275972)

I am interested in a change that keeps the verifier correct and useful (that is, not allowing operations we cannot handle elsewhere), not in the change with the minimal number of LoC. This change should be included in the previous commit instead, I suppose.

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

538

mlir uses camelBack

[limo1996]

2 of ftynse's comments addressed

[limo1996]

getSymbolSource: llvm::Optional instead of -1

[limo1996]

evolve

[ftynse]

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

94

This does not look correct. Assume the concatenated map looks like (d0,d1,d2)[s0] -> (d2,d0,d1) (the symbol is unused for simplicity). The expected ranges should be the same as produced by the inversion method, i.e. (0..allViewSizes[1], 0..allViewSizes[2], 0..allViewSizes[0]), but this code will produce (0..allViewSizes[0], 0..allViewSizes[1], 0..allViewSizes[2]).

I think you need allViewSizes[ <index-of-result-in-map> ] instead of allViewSizes[d.getPosition()], but please double check.

99

Could you please add a comment that explains the math behind this? Affine expressions are not super-easy to follow.

102

Shouldn't this also check that LHS and RHS are binary expressions of a specific kind, like "add" or "floordiv"?

Like: "if the expression has the shape di + dj - sk floordiv CST, the bounds are ... because ... ".

Also explain that it always overwrite the existing range because the range is always smaller in this case. Imagine having a map like (d0,d1)[s0] -> (d0,d1,d0+d1-s0 floordiv 3), you will have computed the bounds for both loops when you hit this case.

<disregard for now>
Normally, we should have computed the maximum across all found lower bounds and the minimum across all found upper bounds.
</disregard for now>

534

Since the parent commit introduced symbol_source as an ODS attribute, it should be possible to just call linalgOp.symbol_source() here and avoid using a hardcoded name.

540–541

unsigned diff = map.getNumSymbols(); ? Also, numSymb sounds like a better name for this?

547

I'm not convinced diff * symbolSource + idx is the correct expression. You actually want the sizes that correspond to symbolSource's operand, which should be computed as the sum of ranks of all operands before symbolSource. (diff is the current number of symbols).

[limo1996]

All comments resolved

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

94

Yes. I moved the logic from the next commit here..

534

This code was removed and simplified so no need for that

540–541

This code was also removed

547

Yes this was also fixed in the previous commit and evolved here..

[limo1996]

All comments resolved

[limo1996]

comment moved

[ftynse]

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

99

The comment I see (m + n - s floordiv 2) is overly concise to qualify as explanation. I was expecting something like: "if the access pattern is (m,n)[s] -> (m + n - s floordiv C)", then the bounds are s floordiv C <= n <= max m - s floordiv C because reason1, reason2.

Actually, I seem to have said the exact same thing in the comment below, which is marked "done". Did you forget to reupload the diff after changes?

102

I don't see checks for binary expression kinds on lhs and rhs. In the current state, this can match something like d0-d1*s0 floordiv 42, in which case the computed bounds will be wrong. Did you forget to upload a new version?

[limo1996]

Evolved changes from parent commit which included changes to emitting loop ranges with symbols

[limo1996]

last comment of nicolas resolved

[ftynse]

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

136

Let's be more explicit: "Reserve is mandatory to avoid a potential undefined behavior with pushing back to smallvector from itself."

Also, please terminate sentences with a dot in comments. Always.

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

99

Thanks for providing more details!

Unfortunately, this confirmed my understanding that the code went too far with eager simplification and produces off-by-one results.

Given

0 <= n < A
0 <= n + m - s floordiv 2 < B

the correct upper bound on m is

m < B - max n + s floordiv 2

We can substitute max n with (A-1) from the first inequality, thus obtaining

m < B - A + 1 + s floordiv 2

Now, the code makes the assumption that s = A, which can be reasonable for the use case. This gives

m < B - A + 1 + A floordiv 2

and then it tries to simplify to

m < B - A floordiv 2

which only holds when A = 1 mod 2, because in this case A = 2 * (A floordiv 2) + 1 and cancels out the +/-1. Otherwise, it should have been m < B - A floordiv 2 + 1. This does not produce out-of-bounds accesses, but ignores the last iteration for even values of A.

Let's not be too eager with simplification and just compute the upper bound on m as m < B - A + 1 + s floordiv 2 without any assumption on the divisibility of A by 2.

102

I still don't see the checks I expected. To be precise,
auto lhs = binOp.getLHS().dyn_cast<AffineBinaryOpExpr>(); does _not_ guarantee that lhs is a sum. I expect to see
if (lhs.getKind != AffineExprKind::Add) continue; and similarly for rhs.

[limo1996]

All comments of Alex resolved

[limo1996]

synced

[ftynse]

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

123

Actually, it shouldn't matter anymore whether you do floordiv 2 or floordiv any-other-value.

125

This could use a better name. Please try really hard to avoid short semantically meaningless names. I recall making this same comment. The purpose of the name may seem obvious today for you, but it won't be when you'll have to debug this code two month from now. Does c2 stand for c * 2, c squared? Given that there is only one use of this variable, just fold it to its use if (minusOne.getValue() != 1) continue;.

142

Can't this be done as part of affine apply? viewSizes is at least partially contained by values. Also, nobody prevents you from adding more inputs to one of affine apply's.

[limo1996]

comments of Alex incorporated

[limo1996]

All comments of Alex done

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

if (attr) {
     unsigned numSymb = indexingMaps[0].getNumSymbols();
     for (unsigned idx = 0; idx < attr.getInt(); idx++)
       symbolsPos += ranks[idx];
 
     // Append or rewrite the end of the value list that corresponds to the
     // values mapping to symbols. Since inside concatinated map symbols are
     // repeated we have to repeat the sizes as well.
     res.resize(concatMap.getNumInputs());
     for (unsigned idx = 0, s = indexingMaps.size(); idx < s; ++idx)
       for (unsigned idx2 = 0; idx2 < numSymb; ++idx2)
         res[numDims + idx * numSymb + idx2] = res[symbolsPos + idx2];
   }
   return res;
 }
 
 /// Returns the values obtained by applying `map` to the list of values.
 /// When non-null, the optional pointer `folder` is used to call into the
 /// `createAndFold` builder method. If `folder` is null, the regular `create`
 /// method is called.
 SmallVector<Value, 4> applyMapToValues(OpBuilder &b, Location loc,
-                                       AffineMap map, ArrayRef<Value> values,
+                                       AffineMap map, ValueRange values,
                                        OperationFolder *folder = nullptr);
 
 /// Returns all the operands of `linalgOp` that are not views.
 /// Asserts that these operands are value types to allow transformations like
 /// tiling to just use the values when cloning `linalgOp`.
 SmallVector<Value, 4> getAssumedNonViewOperands(LinalgOp linalgOp);
 
 /// Apply the permutation defined by `permutation` to `inVec`.

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

emitLoopRanges(OpBuilder &b, Location loc, AffineMap map,
   SmallVector<SubViewOp::Range, 4> res;
   for (unsigned idx = 0, e = map.getNumResults(); idx < e; ++idx) {
     res.push_back(SubViewOp::Range{std_constant_index(0), sizes[idx],
                                    std_constant_index(1)});
   }
   return res;
 }
 
+/// Creates a number of ranges equal to the number of dimensions in the `map`.
+/// The function supports for now only limited number of expressions inside
+/// map results. It expects a non-inverted, concatenated map and last values in
+/// allViewSizes will be applied to the symbols in the map.
+static SmallVector<SubViewOp::Range, 4>
+emitLoopRangesWithSymbols(OpBuilder &b, Location loc, AffineMap map,
+                          ValueRange allViewSizes) {
+  assert(allViewSizes.size() == map.getNumInputs() &&
+         "Number of provided values must match number of inputs to the map.");
+
+  SmallVector<SubViewOp::Range, 4> res(map.getNumDims());
+  for (unsigned idx = 0, e = map.getNumResults(); idx < e; ++idx) {
+    auto result = map.getResult(idx);
+    if (auto d = result.dyn_cast<AffineDimExpr>()) {
+      if (res[d.getPosition()].offset)
+        continue;
+      res[d.getPosition()] = SubViewOp::Range{
+          std_constant_index(0), allViewSizes[idx], std_constant_index(1)};
+    }
+
+    // m + n - s floordiv 2
+    if (auto binOp = result.dyn_cast<AffineBinaryOpExpr>()) {
+      auto lhs = binOp.getLHS().dyn_cast<AffineBinaryOpExpr>();
+      auto rhs = binOp.getRHS().dyn_cast<AffineBinaryOpExpr>();
+      if (!lhs || !rhs)
+        continue;
+
+      auto m = lhs.getLHS().dyn_cast<AffineDimExpr>();
+      auto n = lhs.getRHS().dyn_cast<AffineDimExpr>();
+      auto fDiv = rhs.getLHS().dyn_cast<AffineBinaryOpExpr>();
+      if (!m || !n || !fDiv || fDiv.getKind() != AffineExprKind::FloorDiv ||
+          fDiv.getLHS().getKind() != AffineExprKind::SymbolId ||
+          fDiv.getRHS().getKind() != AffineExprKind::Constant)
+        continue;
+
+      int mPos = m.getPosition();
+      AffineMap fromMap =
+          AffineMap::get(map.getNumDims(), map.getNumSymbols(), fDiv);
+      Value from = applyMapToValues(b, loc, fromMap, allViewSizes).front();
+      auto to = b.create<SubIOp>(loc, allViewSizes[mPos], from);
+      res[mPos] = SubViewOp::Range{from, to, std_constant_index(1)};
+    }
+  }
+  return res;
+}
+
 template <typename IndexedValueType, typename OpType>
 static void inlineRegionAndEmitStore(OpType op, ArrayRef<Value> indexedValues,
                                      ArrayRef<SmallVector<Value, 8>> indexing,
                                      ArrayRef<Value> outputBuffers) {
   assert(op.getOperation()->getNumRegions() == 1 &&
          "Expected single region op");
   auto &b = ScopedContext::getBuilderRef();
   auto &block = op.region().front();
   BlockAndValueMapping map;
   map.map(block.getArguments(), indexedValues);
   for (auto &op : block.without_terminator()) {
     assert(op.getNumRegions() == 0 && "expected a non-nested region");
     auto *newOp = b.clone(op, map);
     map.map(op.getResults(), newOp->getResults());
   }
 
   Operation &terminator = block.back();
   assert(isa<YieldOp>(terminator) &&
          "expected a yield op in the end of the region");
   for (unsigned i = 0, e = terminator.getNumOperands(); i < e; ++i) {
     IndexedValueType O(outputBuffers[i]);
     O(indexing[i]) = map.lookupOrDefault(terminator.getOperand(i));
   }
 }
 
 // Returns a pair that contains input indices and output indices of a
Optional<LinalgLoops> linalgOpToLoopsImpl(Operation *op, OpBuilder &builder) {
   // The flattened loopToOperandRangesMaps is expected to be an invertible
   // permutation map (which is asserted in the inverse calculation).
   auto linalgOp = cast<ConcreteOpTy>(op);
   assert(linalgOp.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   SmallVector<AffineMap, 8> maps = getIndexingMaps(linalgOp);
   SmallVector<Value, 8> sizes = getViewSizes(builder, linalgOp);
   AffineMap map = concatAffineMaps(maps);
-  if (map.getNumSymbols()) {
-    // Ignore symbols for now as they are not supported by inversePermutation.
-    unsigned dims = map.getNumDims();
-    SmallVector<AffineExpr, 8> res;
-    for (auto result : map.getResults())
-      if (auto d = result.dyn_cast<AffineDimExpr>())
-        res.push_back(d);
-
-    map = AffineMap::get(dims, 0, res, map.getContext());
-
-    // Cut off values that would have been applied to symbols
-    sizes.resize(res.size());
-  }
+  SmallVector<SubViewOp::Range, 4> loopRanges;
 
+  if (map.getNumSymbols()) {
+    loopRanges = emitLoopRangesWithSymbols(scope.getBuilderRef(),
+                                           scope.getLocation(), map, sizes);
+  } else {
     AffineMap invertedMap = inversePermutation(map);
     if (!invertedMap)
       return {};
     if (invertedMap.isEmpty()) {
       emitScalarImplementation<IndexedValueTy>({}, linalgOp);
       return LinalgLoops();
     }
 
-  SmallVector<Value, 4> allIvs;
-  auto loopRanges = emitLoopRanges(scope.getBuilderRef(), scope.getLocation(),
-                                   invertedMap, sizes);
+    loopRanges = emitLoopRanges(scope.getBuilderRef(), scope.getLocation(),
+  }
+  SmallVector<Value, 4> allIvs;
   GenerateLoopNest<LoopTy>::doit(
       loopRanges, linalgOp.iterator_types().getValue(), [&](ValueRange ivs) {
         allIvs.append(ivs.begin(), ivs.end());
         emitScalarImplementation<IndexedValueTy>(allIvs, linalgOp);
       });
   // Number of loop ops might be different from the number of ivs since some
   // loops like affine.parallel and scf.parallel have multiple ivs.
   llvm::SetVector<Operation *> loopSet;
   for (Value iv : allIvs) {
     if (!iv)
       return {};
     // The induction variable is a block argument of the entry block of the
     // loop operation.
     BlockArgument ivVal = iv.dyn_cast<BlockArgument>();
     if (!ivVal)
       return {};
     loopSet.insert(ivVal.getOwner()->getParentOp());
   }
   LinalgLoops loops(loopSet.begin(), loopSet.end());
   return loops;
 }
 
 template <typename LoopType, typename ConcreteOp>
 class LinalgRewritePattern : public RewritePattern {
 public:
   explicit LinalgRewritePattern(MLIRContext *context)
       : RewritePattern(ConcreteOp::getOperationName(), 1, context) {}

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

RegionMatcher::matchAsScalarBinaryOp(GenericOp op) {
   if (addPattern.match(&ops.back()))
     return BinaryOpKind::IAdd;
 
   return llvm::None;
 }
 
 static Value emitOrFoldComposedAffineApply(OpBuilder &b, Location loc,
                                            AffineMap map,
-                                           ArrayRef<Value> operandsRef,
+                                           ValueRange operandsRef,
                                            OperationFolder *folder) {
   SmallVector<Value, 4> operands(operandsRef.begin(), operandsRef.end());
   fullyComposeAffineMapAndOperands(&map, &operands);
   canonicalizeMapAndOperands(&map, &operands);
   return folder ? folder->create<AffineApplyOp>(b, loc, map, operands)
                 : b.create<AffineApplyOp>(loc, map, operands);
 }
 
 SmallVector<Value, 4> mlir::linalg::applyMapToValues(OpBuilder &b, Location loc,
                                                      AffineMap map,
-                                                     ArrayRef<Value> values,
+                                                     ValueRange values,
                                                      OperationFolder *folder) {
   SmallVector<Value, 4> res;
   res.reserve(map.getNumResults());
-  unsigned numDims = map.getNumDims();
+  unsigned numDims = map.getNumDims(), numSym = map.getNumSymbols();
   // For each `expr` in `map`, applies the `expr` to the values extracted from
   // ranges. If the resulting application can be folded into a Value, the
   // folding occurs eagerly. Otherwise, an affine.apply operation is emitted.
   for (auto expr : map.getResults()) {
-    AffineMap map = AffineMap::get(numDims, 0, expr);
+    AffineMap map = AffineMap::get(numDims, numSym, expr);
     res.push_back(emitOrFoldComposedAffineApply(b, loc, map, values, folder));
   }
   return res;
 }
 
 /// Returns all the operands of `linalgOp` that are not views.
 /// Asserts that these operands are value types to allow transformations like
 /// tiling to just use the values when cloning `linalgOp`.

mlir/test/Dialect/Linalg/loops.mlir

 // RUN: mlir-opt %s -convert-linalg-to-loops | FileCheck --check-prefix=CHECKLOOP %s
 // RUN: mlir-opt %s -convert-linalg-to-parallel-loops | FileCheck --check-prefix=CHECKPARALLEL %s
 
 // Test that we can lower all the way to LLVM without crashing, don't check results here.
 // RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm -o=/dev/null 2>&1
 
 // CHECKLOOP-DAG: #[[$strided1D:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
 // CHECKLOOP-DAG: #[[$strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 // CHECKLOOP-DAG: #[[$strided3D:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>
 // CHECKLOOP-DAG: #[[$strided4D:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3)>
 // CHECKLOOP-DAG: #[[$clampMinMap:.*]] = affine_map<(d0) -> (d0, 0)>
 
 // CHECKLOOP-DAG: #[[$stride1Dilation1:.*]] = affine_map<(d0, d1) -> (d0  + d1)>
 // CHECKLOOP-DAG: #[[$stride2Dilation1:.*]] = affine_map<(d0, d1) -> (d0 * 2 + d1)>
 // CHECKLOOP-DAG: #[[$stride2Dilation4:.*]] = affine_map<(d0, d1) -> (d0 * 2 + d1 * 4)>
 // CHECKLOOP-DAG: #[[$stride3Dilation5:.*]] = affine_map<(d0, d1) -> (d0 * 3 + d1 * 5)>
+// CHECKLOOP-DAG: #[[$convHalf:.*]] = affine_map<()[s0] -> (s0 floordiv 2)>
 // CHECKLOOP-DAG: #[[$convMap:.*]] = affine_map<(d0, d1)[s0] -> (d0 + d1 - s0 floordiv 2)>
 
 // CHECKPARALLEL-DAG: #[[$strided1D:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
 // CHECKPARALLEL-DAG: #[[$strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 // CHECKPARALLEL-DAG: #[[$strided3D:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>
 // CHECKPARALLEL-DAG: #[[$strided4D:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3)>
 // CHECKPARALLEL-DAG: #[[$clampMinMap:.*]] = affine_map<(d0) -> (d0, 0)>
 
 // CHECKPARALLEL-DAG: #[[$stride1Dilation1:.*]] = affine_map<(d0, d1) -> (d0  + d1)>
 // CHECKPARALLEL-DAG: #[[$stride2Dilation1:.*]] = affine_map<(d0, d1) -> (d0 * 2 + d1)>
 // CHECKPARALLEL-DAG: #[[$stride2Dilation4:.*]] = affine_map<(d0, d1) -> (d0 * 2 + d1 * 4)>
 // CHECKPARALLEL-DAG: #[[$stride3Dilation5:.*]] = affine_map<(d0, d1) -> (d0 * 3 + d1 * 5)>
+// CHECKPARALLEL-DAG: #[[$convHalf:.*]] = affine_map<()[s0] -> (s0 floordiv 2)>
 // CHECKPARALLEL-DAG: #[[$convMap:.*]] = affine_map<(d0, d1)[s0] -> (d0 + d1 - s0 floordiv 2)>
 
 
 func @matmul(%arg0: memref<?xi8>, %M: index, %N: index, %K: index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %A = view %arg0[%c0][%M, %K] : memref<?xi8> to memref<?x?xf32>
   %B = view %arg0[%c0][%K, %N] : memref<?xi8> to memref<?x?xf32>
func @conv1d(%filter : memref<?xf32>, %in : memref<?xf32>, %out :  memref<?xf32>) -> () {
   return
 }
 
 // CHECKLOOP-LABEL: @conv1d
 //  CHECKLOOP-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?xf32>
 //  CHECKLOOP-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?xf32>
 //  CHECKLOOP-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?xf32>
 //       CHECKLOOP: %[[c0:.*]] = constant 0 : index
-//       CHECKLOOP: %[[dim0:.*]] = dim %[[arg0]], %[[c0]] : memref<?xf32>
-//       CHECKLOOP: %[[dim1:.*]] = dim %[[arg2]], %[[c0]] : memref<?xf32>
-//       CHECKLOOP: scf.for %[[b:.*]] = %{{.*}} to %[[dim1]] step %{{.*}} {
-//       CHECKLOOP:   scf.for %[[m:.*]] = %{{.*}} to %[[dim0]] step %{{.*}} {
+//       CHECKLOOP: %[[dim0:.*]] = dim %[[arg1]], %[[c0]] : memref<?xf32>
+//       CHECKLOOP: %[[dim1:.*]] = dim %[[arg0]], %[[c0]] : memref<?xf32>
+//       CHECKLOOP: %[[half:.*]] = affine.apply #[[$convHalf]]()[%[[dim1]]]
+//       CHECKLOOP: %[[sizeMinusHalf:.*]] = subi %[[dim0]], %[[half]] : index
+//       CHECKLOOP: scf.for %[[b:.*]] = %[[half]] to %[[sizeMinusHalf]] step %{{.*}} {
+//       CHECKLOOP:   scf.for %[[m:.*]] = %{{.*}} to %[[dim1]] step %{{.*}} {
 //       CHECKLOOP:     %[[dim2:.*]] = dim %[[arg0]], %[[c0]] : memref<?xf32>
 //       CHECKLOOP:     %[[aff:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim2]]]
 //       CHECKLOOP:     %[[va:.*]] = load %[[arg1]][%[[aff]]] : memref<?xf32>
 //       CHECKLOOP:     %[[vb:.*]] = load %[[arg0]][%[[m]]] : memref<?xf32>
 //       CHECKLOOP:     %[[vc:.*]] = load %[[arg2]][%[[b]]] : memref<?xf32>
 //       CHECKLOOP:     %[[inc:.*]] = mulf %[[va]], %[[vb]] : f32
 //       CHECKLOOP:     %[[res:.*]] = addf %[[vc]], %[[inc]] : f32
 //       CHECKLOOP:     store %[[res]], %[[arg2]][%[[b]]] : memref<?xf32>
 
 // CHECKPARALLEL-LABEL: @conv1d
 //  CHECKPARALLEL-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?xf32>
 //  CHECKPARALLEL-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?xf32>
 //  CHECKPARALLEL-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?xf32>
 //       CHECKPARALLEL: %[[c0:.*]] = constant 0 : index
-//       CHECKPARALLEL: %[[dim0:.*]] = dim %[[arg0]], %[[c0]] : memref<?xf32>
-//       CHECKPARALLEL: %[[dim1:.*]] = dim %[[arg2]], %[[c0]] : memref<?xf32>
-//       CHECKPARALLEL: scf.parallel (%[[b:.*]], %[[m:.*]]) = (%{{.*}}, %{{.*}}) to (%[[dim1]], %[[dim0]]) step ({{.*}}) {
+//       CHECKPARALLEL: %[[dim0:.*]] = dim %[[arg1]], %[[c0]] : memref<?xf32>
+//       CHECKPARALLEL: %[[dim1:.*]] = dim %[[arg0]], %[[c0]] : memref<?xf32>
+//       CHECKPARALLEL: %[[half:.*]] = affine.apply #[[$convHalf]]()[%[[dim1]]]
+//       CHECKPARALLEL: %[[sizeMinusHalf:.*]] = subi %[[dim0]], %[[half]] : index
+//       CHECKPARALLEL: scf.parallel (%[[b:.*]], %[[m:.*]]) = (%[[half]], %{{.*}}) to (%[[sizeMinusHalf]], %[[dim1]]) step ({{.*}}) {
 //       CHECKPARALLEL:   %[[dim2:.*]] = dim %[[arg0]], %[[c0]] : memref<?xf32>
 //       CHECKPARALLEL:   %[[aff:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim2]]]
 //       CHECKPARALLEL:   %[[va:.*]] = load %[[arg1]][%[[aff]]] : memref<?xf32>
 //       CHECKPARALLEL:   %[[vb:.*]] = load %[[arg0]][%[[m]]] : memref<?xf32>
 //       CHECKPARALLEL:   %[[vc:.*]] = load %[[arg2]][%[[b]]] : memref<?xf32>
 //       CHECKPARALLEL:   %[[inc:.*]] = mulf %[[va]], %[[vb]] : f32
 //       CHECKPARALLEL:   %[[res:.*]] = addf %[[vc]], %[[inc]] : f32
 //       CHECKPARALLEL:   store %[[res]], %[[arg2]][%[[b]]] : memref<?xf32>

 }
 
 // CHECKLOOP-LABEL: @conv2d
 //  CHECKLOOP-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?x?xf32>
 //  CHECKLOOP-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?x?xf32>
 //  CHECKLOOP-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?x?xf32>
 //       CHECKLOOP: %[[c0:.*]] = constant 0 : index
 //       CHECKLOOP: %[[c1:.*]] = constant 1 : index
-//       CHECKLOOP: %[[dim0:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?xf32>
-//       CHECKLOOP: %[[dim1:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?xf32>
-//       CHECKLOOP: %[[dim2:.*]] = dim %[[arg2]], %[[c0]] : memref<?x?xf32>
-//       CHECKLOOP: %[[dim3:.*]] = dim %[[arg2]], %[[c1]] : memref<?x?xf32>
-//       CHECKLOOP: scf.for %[[i0:.*]] = %{{.*}} to %[[dim2]] step %{{.*}} {
-//       CHECKLOOP:   scf.for %[[i1:.*]] = %{{.*}} to %[[dim3]] step %{{.*}} {
-//       CHECKLOOP:     scf.for %[[i2:.*]] = %{{.*}} to %[[dim0]] step %{{.*}} {
-//       CHECKLOOP:       scf.for %[[i3:.*]] = %{{.*}} to %[[dim1]] step %{{.*}} {
+//       CHECKLOOP: %[[dim0:.*]] = dim %[[arg1]], %[[c0]] : memref<?x?xf32>
+//       CHECKLOOP: %[[dim1:.*]] = dim %[[arg1]], %[[c1]] : memref<?x?xf32>
+//       CHECKLOOP: %[[dim2:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?xf32>
+//       CHECKLOOP: %[[dim3:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?xf32>
+//       CHECKLOOP: %[[half1:.*]] = affine.apply #[[$convHalf]]()[%[[dim2]]]
+//       CHECKLOOP: %[[sizeMinusHalf1:.*]] = subi %[[dim0]], %[[half1]] : index
+//       CHECKLOOP: %[[half2:.*]] = affine.apply #[[$convHalf]]()[%[[dim3]]]
+//       CHECKLOOP: %[[sizeMinusHalf2:.*]] = subi %[[dim1]], %[[half2]] : index
+//       CHECKLOOP: scf.for %[[i0:.*]] = %[[half1]] to %[[sizeMinusHalf1]] step %{{.*}} {
+//       CHECKLOOP:   scf.for %[[i1:.*]] = %[[half2]] to %[[sizeMinusHalf2]] step %{{.*}} {
+//       CHECKLOOP:     scf.for %[[i2:.*]] = %{{.*}} to %[[dim2]] step %{{.*}} {
+//       CHECKLOOP:       scf.for %[[i3:.*]] = %{{.*}} to %[[dim3]] step %{{.*}} {
 //       CHECKLOOP:         %[[dim4:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?xf32>
 //       CHECKLOOP:         %[[dim5:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?xf32>
 //       CHECKLOOP:         %[[aff1:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim4]]]
 //       CHECKLOOP:         %[[aff2:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim5]]]
 //       CHECKLOOP:         %[[va:.*]] = load %[[arg1]][%[[aff1]], %[[aff2]]] : memref<?x?xf32>
 //       CHECKLOOP:         %[[vb:.*]] = load %[[arg0]][%[[i2]], %[[i3]]] : memref<?x?xf32>
 //       CHECKLOOP:         %[[vc:.*]] = load %[[arg2]][%[[i0]], %[[i1]]] : memref<?x?xf32>
 //       CHECKLOOP:         %[[inc:.*]] = mulf %[[va]], %[[vb]] : f32
 //       CHECKLOOP:         %[[res:.*]] = addf %[[vc]], %[[inc]] : f32
 //       CHECKLOOP:         store %[[res]], %[[arg2]][%[[i0]], %[[i1]]] : memref<?x?xf32>
 
 // CHECKPARALLEL-LABEL: @conv2d
 //  CHECKPARALLEL-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?x?xf32>
 //  CHECKPARALLEL-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?x?xf32>
 //  CHECKPARALLEL-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?x?xf32>
 //       CHECKPARALLEL: %[[c0:.*]] = constant 0 : index
 //       CHECKPARALLEL: %[[c1:.*]] = constant 1 : index
-//       CHECKPARALLEL: %[[dim0:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?xf32>
-//       CHECKPARALLEL: %[[dim1:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?xf32>
-//       CHECKPARALLEL: %[[dim2:.*]] = dim %[[arg2]], %[[c0]] : memref<?x?xf32>
-//       CHECKPARALLEL: %[[dim3:.*]] = dim %[[arg2]], %[[c1]] : memref<?x?xf32>
-//       CHECKPARALLEL: scf.parallel (%[[i0:.*]], %[[i1:.*]], %[[i2:.*]], %[[i3:.*]]) = (%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) to (%[[dim2]], %[[dim3]], %[[dim0]], %[[dim1]]) step ({{.*}}) {
+//       CHECKPARALLEL: %[[dim0:.*]] = dim %[[arg1]], %[[c0]] : memref<?x?xf32>
+//       CHECKPARALLEL: %[[dim1:.*]] = dim %[[arg1]], %[[c1]] : memref<?x?xf32>
+//       CHECKPARALLEL: %[[dim2:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?xf32>
+//       CHECKPARALLEL: %[[dim3:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?xf32>
+//       CHECKPARALLEL: %[[half1:.*]] = affine.apply #[[$convHalf]]()[%[[dim2]]]
+//       CHECKPARALLEL: %[[sizeMinusHalf1:.*]] = subi %[[dim0]], %[[half1]] : index
+//       CHECKPARALLEL: %[[half2:.*]] = affine.apply #[[$convHalf]]()[%[[dim3]]]
+//       CHECKPARALLEL: %[[sizeMinusHalf2:.*]] = subi %[[dim1]], %[[half2]] : index
+//       CHECKPARALLEL: scf.parallel (%[[i0:.*]], %[[i1:.*]], %[[i2:.*]], %[[i3:.*]]) = (%[[half1]], %[[half2]], %{{.*}}, %{{.*}}) to (%[[sizeMinusHalf1]], %[[sizeMinusHalf2]], %[[dim2]], %[[dim3]]) step ({{.*}}) {
 //       CHECKPARALLEL:   %[[dim4:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?xf32>
 //       CHECKPARALLEL:   %[[dim5:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?xf32>
 //       CHECKPARALLEL:   %[[aff1:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim4]]]
 //       CHECKPARALLEL:   %[[aff2:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim5]]]
 //       CHECKPARALLEL:   %[[va:.*]] = load %[[arg1]][%[[aff1]], %[[aff2]]] : memref<?x?xf32>
 //       CHECKPARALLEL:   %[[vb:.*]] = load %[[arg0]][%[[i2]], %[[i3]]] : memref<?x?xf32>
 //       CHECKPARALLEL:   %[[vc:.*]] = load %[[arg2]][%[[i0]], %[[i1]]] : memref<?x?xf32>
 //       CHECKPARALLEL:   %[[inc:.*]] = mulf %[[va]], %[[vb]] : f32

 
 // CHECKLOOP-LABEL: @conv3d
 //  CHECKLOOP-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
 //  CHECKLOOP-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
 //  CHECKLOOP-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
 //       CHECKLOOP: %[[c0:.*]] = constant 0 : index
 //       CHECKLOOP: %[[c1:.*]] = constant 1 : index
 //       CHECKLOOP: %[[c2:.*]] = constant 2 : index
-//       CHECKLOOP: %[[dim0:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?x?xf32>
-//       CHECKLOOP: %[[dim1:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?x?xf32>
-//       CHECKLOOP: %[[dim2:.*]] = dim %[[arg0]], %[[c2]] : memref<?x?x?xf32>
-//       CHECKLOOP: %[[dim3:.*]] = dim %[[arg2]], %[[c0]] : memref<?x?x?xf32>
-//       CHECKLOOP: %[[dim4:.*]] = dim %[[arg2]], %[[c1]] : memref<?x?x?xf32>
-//       CHECKLOOP: %[[dim5:.*]] = dim %[[arg2]], %[[c2]] : memref<?x?x?xf32>
-//       CHECKLOOP: scf.for %[[i0:.*]] = %{{.*}} to %[[dim3]] step %{{.*}} {
-//       CHECKLOOP:   scf.for %[[i1:.*]] = %{{.*}} to %[[dim4]] step %{{.*}} {
-//       CHECKLOOP:     scf.for %[[i2:.*]] = %{{.*}} to %[[dim5]] step %{{.*}} {
-//       CHECKLOOP:       scf.for %[[i3:.*]] = %{{.*}} to %[[dim0]] step %{{.*}} {
-//       CHECKLOOP:         scf.for %[[i4:.*]] = %{{.*}} to %[[dim1]] step %{{.*}} {
-//       CHECKLOOP:           scf.for %[[i5:.*]] = %{{.*}} to %[[dim2]] step %{{.*}} {
+//       CHECKLOOP: %[[dim0:.*]] = dim %[[arg1]], %[[c0]] : memref<?x?x?xf32>
+//       CHECKLOOP: %[[dim1:.*]] = dim %[[arg1]], %[[c1]] : memref<?x?x?xf32>
+//       CHECKLOOP: %[[dim2:.*]] = dim %[[arg1]], %[[c2]] : memref<?x?x?xf32>
+//       CHECKLOOP: %[[dim3:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?x?xf32>
+//       CHECKLOOP: %[[dim4:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?x?xf32>
+//       CHECKLOOP: %[[dim5:.*]] = dim %[[arg0]], %[[c2]] : memref<?x?x?xf32>
+//       CHECKLOOP: %[[half1:.*]] = affine.apply #[[$convHalf]]()[%[[dim3]]]
+//       CHECKLOOP: %[[sizeMinusHalf1:.*]] = subi %[[dim0]], %[[half1]] : index
+//       CHECKLOOP: %[[half2:.*]] = affine.apply #[[$convHalf]]()[%[[dim4]]]
+//       CHECKLOOP: %[[sizeMinusHalf2:.*]] = subi %[[dim1]], %[[half2]] : index
+//       CHECKLOOP: %[[half3:.*]] = affine.apply #[[$convHalf]]()[%[[dim5]]]
+//       CHECKLOOP: %[[sizeMinusHalf3:.*]] = subi %[[dim2]], %[[half3]] : index
+//       CHECKLOOP: scf.for %[[i0:.*]] = %[[half1]] to %[[sizeMinusHalf1]] step %{{.*}} {
+//       CHECKLOOP:   scf.for %[[i1:.*]] = %[[half2]] to %[[sizeMinusHalf2]] step %{{.*}} {
+//       CHECKLOOP:     scf.for %[[i2:.*]] = %[[half3]] to %[[sizeMinusHalf3]] step %{{.*}} {
+//       CHECKLOOP:       scf.for %[[i3:.*]] = %{{.*}} to %[[dim3]] step %{{.*}} {
+//       CHECKLOOP:         scf.for %[[i4:.*]] = %{{.*}} to %[[dim4]] step %{{.*}} {
+//       CHECKLOOP:           scf.for %[[i5:.*]] = %{{.*}} to %[[dim5]] step %{{.*}} {
 //       CHECKLOOP:             %[[dim6:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?x?xf32>
 //       CHECKLOOP:             %[[dim7:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?x?xf32>
 //       CHECKLOOP:             %[[dim8:.*]] = dim %[[arg0]], %[[c2]] : memref<?x?x?xf32>
 //       CHECKLOOP:             %[[aff1:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim6]]]
 //       CHECKLOOP:             %[[aff2:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim7]]]
 //       CHECKLOOP:             %[[aff3:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim8]]]
 //       CHECKLOOP:             %[[va:.*]] = load %[[arg1]][%[[aff1]], %[[aff2]], %[[aff3]]] : memref<?x?x?xf32>
 //       CHECKLOOP:             %[[vb:.*]] = load %[[arg0]][%[[i3]], %[[i4]], %[[i5]]] : memref<?x?x?xf32>
 //       CHECKLOOP:             %[[vc:.*]] = load %[[arg2]][%[[i0]], %[[i1]], %[[i2]]] : memref<?x?x?xf32>
 //       CHECKLOOP:             %[[inc:.*]] = mulf %[[va]], %[[vb]] : f32
 //       CHECKLOOP:             %[[res:.*]] = addf %[[vc]], %[[inc]] : f32
 //       CHECKLOOP:             store %[[res]], %[[arg2]][%[[i0]], %[[i1]], %[[i2]]] : memref<?x?x?xf32>
 
 // CHECKPARALLEL-LABEL: @conv3d
 //  CHECKPARALLEL-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
 //  CHECKPARALLEL-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
 //  CHECKPARALLEL-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
 //       CHECKPARALLEL: %[[c0:.*]] = constant 0 : index
 //       CHECKPARALLEL: %[[c1:.*]] = constant 1 : index
 //       CHECKPARALLEL: %[[c2:.*]] = constant 2 : index
-//       CHECKPARALLEL: %[[dim0:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?x?xf32>
-//       CHECKPARALLEL: %[[dim1:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?x?xf32>
-//       CHECKPARALLEL: %[[dim2:.*]] = dim %[[arg0]], %[[c2]] : memref<?x?x?xf32>
-//       CHECKPARALLEL: %[[dim3:.*]] = dim %[[arg2]], %[[c0]] : memref<?x?x?xf32>
-//       CHECKPARALLEL: %[[dim4:.*]] = dim %[[arg2]], %[[c1]] : memref<?x?x?xf32>
-//       CHECKPARALLEL: %[[dim5:.*]] = dim %[[arg2]], %[[c2]] : memref<?x?x?xf32>
-//       CHECKPARALLEL: scf.parallel (%[[i0:.*]], %[[i1:.*]], %[[i2:.*]], %[[i3:.*]], %[[i4:.*]], %[[i5:.*]]) = (%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) to (%[[dim3]], %[[dim4]], %[[dim5]], %[[dim0]], %[[dim1]], %[[dim2]]) step ({{.*}}) {
+//       CHECKPARALLEL: %[[dim0:.*]] = dim %[[arg1]], %[[c0]] : memref<?x?x?xf32>
+//       CHECKPARALLEL: %[[dim1:.*]] = dim %[[arg1]], %[[c1]] : memref<?x?x?xf32>
+//       CHECKPARALLEL: %[[dim2:.*]] = dim %[[arg1]], %[[c2]] : memref<?x?x?xf32>
+//       CHECKPARALLEL: %[[dim3:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?x?xf32>
+//       CHECKPARALLEL: %[[dim4:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?x?xf32>
+//       CHECKPARALLEL: %[[dim5:.*]] = dim %[[arg0]], %[[c2]] : memref<?x?x?xf32>
+//       CHECKPARALLEL: %[[half1:.*]] = affine.apply #[[$convHalf]]()[%[[dim3]]]
+//       CHECKPARALLEL: %[[sizeMinusHalf1:.*]] = subi %[[dim0]], %[[half1]] : index
+//       CHECKPARALLEL: %[[half2:.*]] = affine.apply #[[$convHalf]]()[%[[dim4]]]
+//       CHECKPARALLEL: %[[sizeMinusHalf2:.*]] = subi %[[dim1]], %[[half2]] : index
+//       CHECKPARALLEL: %[[half3:.*]] = affine.apply #[[$convHalf]]()[%[[dim5]]]
+//       CHECKPARALLEL: %[[sizeMinusHalf3:.*]] = subi %[[dim2]], %[[half3]] : index
+//       CHECKPARALLEL: scf.parallel (%[[i0:.*]], %[[i1:.*]], %[[i2:.*]], %[[i3:.*]], %[[i4:.*]], %[[i5:.*]]) = (%[[half1]], %[[half2]], %[[half3]], %{{.*}}, %{{.*}}, %{{.*}}) to (%[[sizeMinusHalf1]], %[[sizeMinusHalf2]], %[[sizeMinusHalf3]], %[[dim3]], %[[dim4]], %[[dim5]]) step ({{.*}}) {
 //       CHECKPARALLEL:   %[[dim6:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?x?xf32>
 //       CHECKPARALLEL:   %[[dim7:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?x?xf32>
 //       CHECKPARALLEL:   %[[dim8:.*]] = dim %[[arg0]], %[[c2]] : memref<?x?x?xf32>
 //       CHECKPARALLEL:   %[[aff1:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim6]]]
 //       CHECKPARALLEL:   %[[aff2:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim7]]]
 //       CHECKPARALLEL:   %[[aff3:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim8]]]
 //       CHECKPARALLEL:   %[[va:.*]] = load %[[arg1]][%[[aff1]], %[[aff2]], %[[aff3]]] : memref<?x?x?xf32>
 //       CHECKPARALLEL:   %[[vb:.*]] = load %[[arg0]][%[[i3]], %[[i4]], %[[i5]]] : memref<?x?x?xf32>

 // CHECKLOOP-LABEL: @conv4d
 //  CHECKLOOP-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32>
 //  CHECKLOOP-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32>
 //  CHECKLOOP-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32>
 //       CHECKLOOP: %[[c0:.*]] = constant 0 : index
 //       CHECKLOOP: %[[c1:.*]] = constant 1 : index
 //       CHECKLOOP: %[[c2:.*]] = constant 2 : index
 //       CHECKLOOP: %[[c3:.*]] = constant 3 : index
-//       CHECKLOOP: %[[dim0:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?x?x?xf32>
-//       CHECKLOOP: %[[dim1:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?x?x?xf32>
-//       CHECKLOOP: %[[dim2:.*]] = dim %[[arg0]], %[[c2]] : memref<?x?x?x?xf32>
-//       CHECKLOOP: %[[dim3:.*]] = dim %[[arg0]], %[[c3]] : memref<?x?x?x?xf32>
-//       CHECKLOOP: %[[dim4:.*]] = dim %[[arg2]], %[[c0]] : memref<?x?x?x?xf32>
-//       CHECKLOOP: %[[dim5:.*]] = dim %[[arg2]], %[[c1]] : memref<?x?x?x?xf32>
-//       CHECKLOOP: %[[dim6:.*]] = dim %[[arg2]], %[[c2]] : memref<?x?x?x?xf32>
-//       CHECKLOOP: %[[dim7:.*]] = dim %[[arg2]], %[[c3]] : memref<?x?x?x?xf32>
-//       CHECKLOOP: scf.for %[[i0:.*]] = %{{.*}} to %[[dim4]] step %{{.*}} {
-//       CHECKLOOP:   scf.for %[[i1:.*]] = %{{.*}} to %[[dim5]] step %{{.*}} {
-//       CHECKLOOP:     scf.for %[[i2:.*]] = %{{.*}} to %[[dim6]] step %{{.*}} {
-//       CHECKLOOP:       scf.for %[[i3:.*]] = %{{.*}} to %[[dim7]] step %{{.*}} {
-//       CHECKLOOP:         scf.for %[[i4:.*]] = %{{.*}} to %[[dim0]] step %{{.*}} {
-//       CHECKLOOP:           scf.for %[[i5:.*]] = %{{.*}} to %[[dim1]] step %{{.*}} {
-//       CHECKLOOP:             scf.for %[[i6:.*]] = %{{.*}} to %[[dim2]] step %{{.*}} {
-//       CHECKLOOP:               scf.for %[[i7:.*]] = %{{.*}} to %[[dim3]] step %{{.*}} {
+//       CHECKLOOP: %[[dim0:.*]] = dim %[[arg1]], %[[c0]] : memref<?x?x?x?xf32>
+//       CHECKLOOP: %[[dim1:.*]] = dim %[[arg1]], %[[c1]] : memref<?x?x?x?xf32>
+//       CHECKLOOP: %[[dim2:.*]] = dim %[[arg1]], %[[c2]] : memref<?x?x?x?xf32>
+//       CHECKLOOP: %[[dim3:.*]] = dim %[[arg1]], %[[c3]] : memref<?x?x?x?xf32>
+//       CHECKLOOP: %[[dim4:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?x?x?xf32>
+//       CHECKLOOP: %[[dim5:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?x?x?xf32>
+//       CHECKLOOP: %[[dim6:.*]] = dim %[[arg0]], %[[c2]] : memref<?x?x?x?xf32>
+//       CHECKLOOP: %[[dim7:.*]] = dim %[[arg0]], %[[c3]] : memref<?x?x?x?xf32>
+//       CHECKLOOP: %[[half1:.*]] = affine.apply #[[$convHalf]]()[%[[dim4]]]
+//       CHECKLOOP: %[[sizeMinusHalf1:.*]] = subi %[[dim0]], %[[half1]] : index
+//       CHECKLOOP: %[[half2:.*]] = affine.apply #[[$convHalf]]()[%[[dim5]]]
+//       CHECKLOOP: %[[sizeMinusHalf2:.*]] = subi %[[dim1]], %[[half2]] : index
+//       CHECKLOOP: %[[half3:.*]] = affine.apply #[[$convHalf]]()[%[[dim6]]]
+//       CHECKLOOP: %[[sizeMinusHalf3:.*]] = subi %[[dim2]], %[[half3]] : index
+//       CHECKLOOP: %[[half4:.*]] = affine.apply #[[$convHalf]]()[%[[dim7]]]
+//       CHECKLOOP: %[[sizeMinusHalf4:.*]] = subi %[[dim3]], %[[half4]] : index
+//       CHECKLOOP: scf.for %[[i0:.*]] = %[[half1]] to %[[sizeMinusHalf1]] step %{{.*}} {
+//       CHECKLOOP:   scf.for %[[i1:.*]] = %[[half2]] to %[[sizeMinusHalf2]] step %{{.*}} {
+//       CHECKLOOP:     scf.for %[[i2:.*]] = %[[half3]] to %[[sizeMinusHalf3]] step %{{.*}} {
+//       CHECKLOOP:       scf.for %[[i3:.*]] = %[[half4]] to %[[sizeMinusHalf4]] step %{{.*}} {
+//       CHECKLOOP:         scf.for %[[i4:.*]] = %{{.*}} to %[[dim4]] step %{{.*}} {
+//       CHECKLOOP:           scf.for %[[i5:.*]] = %{{.*}} to %[[dim5]] step %{{.*}} {
+//       CHECKLOOP:             scf.for %[[i6:.*]] = %{{.*}} to %[[dim6]] step %{{.*}} {
+//       CHECKLOOP:               scf.for %[[i7:.*]] = %{{.*}} to %[[dim7]] step %{{.*}} {
 //       CHECKLOOP:                 %[[dim8:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?x?x?xf32>
 //       CHECKLOOP:                 %[[dim9:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?x?x?xf32>
 //       CHECKLOOP:                 %[[dim10:.*]] = dim %[[arg0]], %[[c2]] : memref<?x?x?x?xf32>
 //       CHECKLOOP:                 %[[dim11:.*]] = dim %[[arg0]], %[[c3]] : memref<?x?x?x?xf32>
 //       CHECKLOOP:                 %[[aff1:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim8]]]
 //       CHECKLOOP:                 %[[aff2:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim9]]]
 //       CHECKLOOP:                 %[[aff3:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim10]]]
 //       CHECKLOOP:                 %[[aff4:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim11]]]
 //       CHECKLOOP:                 %[[va:.*]] = load %[[arg1]][%[[aff1]], %[[aff2]], %[[aff3]], %[[aff4]]] : memref<?x?x?x?xf32>
 //       CHECKLOOP:                 %[[vb:.*]] = load %[[arg0]][%[[i4]], %[[i5]], %[[i6]], %[[i7]]] : memref<?x?x?x?xf32>
 //       CHECKLOOP:                 %[[vc:.*]] = load %[[arg2]][%[[i0]], %[[i1]], %[[i2]], %[[i3]]] : memref<?x?x?x?xf32>
 //       CHECKLOOP:                 %[[inc:.*]] = mulf %[[va]], %[[vb]] : f32
 //       CHECKLOOP:                 %[[res:.*]] = addf %[[vc]], %[[inc]] : f32
 //       CHECKLOOP:                 store %[[res]], %[[arg2]][%[[i0]], %[[i1]], %[[i2]], %[[i3]]] : memref<?x?x?x?xf32>
 
 // CHECKPARALLEL-LABEL: @conv4d
 //  CHECKPARALLEL-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32>
 //  CHECKPARALLEL-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32>
 //  CHECKPARALLEL-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32>
 //       CHECKPARALLEL: %[[c0:.*]] = constant 0 : index
 //       CHECKPARALLEL: %[[c1:.*]] = constant 1 : index
 //       CHECKPARALLEL: %[[c2:.*]] = constant 2 : index
 //       CHECKPARALLEL: %[[c3:.*]] = constant 3 : index
-//       CHECKPARALLEL: %[[dim0:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?x?x?xf32>
-//       CHECKPARALLEL: %[[dim1:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?x?x?xf32>
-//       CHECKPARALLEL: %[[dim2:.*]] = dim %[[arg0]], %[[c2]] : memref<?x?x?x?xf32>
-//       CHECKPARALLEL: %[[dim3:.*]] = dim %[[arg0]], %[[c3]] : memref<?x?x?x?xf32>
-//       CHECKPARALLEL: %[[dim4:.*]] = dim %[[arg2]], %[[c0]] : memref<?x?x?x?xf32>
-//       CHECKPARALLEL: %[[dim5:.*]] = dim %[[arg2]], %[[c1]] : memref<?x?x?x?xf32>
-//       CHECKPARALLEL: %[[dim6:.*]] = dim %[[arg2]], %[[c2]] : memref<?x?x?x?xf32>
-//       CHECKPARALLEL: %[[dim7:.*]] = dim %[[arg2]], %[[c3]] : memref<?x?x?x?xf32>
-//       CHECKPARALLEL: scf.parallel (%[[i0:.*]], %[[i1:.*]], %[[i2:.*]], %[[i3:.*]], %[[i4:.*]], %[[i5:.*]], %[[i6:.*]], %[[i7:.*]]) = (%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) to (%[[dim4]], %[[dim5]], %[[dim6]], %[[dim7]], %[[dim0]], %[[dim1]], %[[dim2]], %[[dim3]]) step ({{.*}}) {
+//       CHECKPARALLEL: %[[dim0:.*]] = dim %[[arg1]], %[[c0]] : memref<?x?x?x?xf32>
+//       CHECKPARALLEL: %[[dim1:.*]] = dim %[[arg1]], %[[c1]] : memref<?x?x?x?xf32>
+//       CHECKPARALLEL: %[[dim2:.*]] = dim %[[arg1]], %[[c2]] : memref<?x?x?x?xf32>
+//       CHECKPARALLEL: %[[dim3:.*]] = dim %[[arg1]], %[[c3]] : memref<?x?x?x?xf32>
+//       CHECKPARALLEL: %[[dim4:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?x?x?xf32>
+//       CHECKPARALLEL: %[[dim5:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?x?x?xf32>
+//       CHECKPARALLEL: %[[dim6:.*]] = dim %[[arg0]], %[[c2]] : memref<?x?x?x?xf32>
+//       CHECKPARALLEL: %[[dim7:.*]] = dim %[[arg0]], %[[c3]] : memref<?x?x?x?xf32>
+//       CHECKPARALLEL: %[[half1:.*]] = affine.apply #[[$convHalf]]()[%[[dim4]]]
+//       CHECKPARALLEL: %[[sizeMinusHalf1:.*]] = subi %[[dim0]], %[[half1]] : index
+//       CHECKPARALLEL: %[[half2:.*]] = affine.apply #[[$convHalf]]()[%[[dim5]]]
+//       CHECKPARALLEL: %[[sizeMinusHalf2:.*]] = subi %[[dim1]], %[[half2]] : index
+//       CHECKPARALLEL: %[[half3:.*]] = affine.apply #[[$convHalf]]()[%[[dim6]]]
+//       CHECKPARALLEL: %[[sizeMinusHalf3:.*]] = subi %[[dim2]], %[[half3]] : index
+//       CHECKPARALLEL: %[[half4:.*]] = affine.apply #[[$convHalf]]()[%[[dim7]]]
+//       CHECKPARALLEL: %[[sizeMinusHalf4:.*]] = subi %[[dim3]], %[[half4]] : index
+//       CHECKPARALLEL: scf.parallel (%[[i0:.*]], %[[i1:.*]], %[[i2:.*]], %[[i3:.*]], %[[i4:.*]], %[[i5:.*]], %[[i6:.*]], %[[i7:.*]]) = (%[[half1]], %[[half2]], %[[half3]], %[[half4]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) to (%[[sizeMinusHalf1]], %[[sizeMinusHalf2]], %[[sizeMinusHalf3]], %[[sizeMinusHalf4]], %[[dim4]], %[[dim5]], %[[dim6]], %[[dim7]]) step ({{.*}}) {
 //       CHECKPARALLEL:   %[[dim8:.*]] = dim %[[arg0]], %[[c0]] : memref<?x?x?x?xf32>
 //       CHECKPARALLEL:   %[[dim9:.*]] = dim %[[arg0]], %[[c1]] : memref<?x?x?x?xf32>
 //       CHECKPARALLEL:   %[[dim10:.*]] = dim %[[arg0]], %[[c2]] : memref<?x?x?x?xf32>
 //       CHECKPARALLEL:   %[[dim11:.*]] = dim %[[arg0]], %[[c3]] : memref<?x?x?x?xf32>
 //       CHECKPARALLEL:   %[[aff1:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim8]]]
 //       CHECKPARALLEL:   %[[aff2:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim9]]]
 //       CHECKPARALLEL:   %[[aff3:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim10]]]
 //       CHECKPARALLEL:   %[[aff4:.*]] = affine.apply #[[$convMap]](%{{.*}}, %{{.*}})[%[[dim11]]]
 //       CHECKPARALLEL:   %[[va:.*]] = load %[[arg1]][%[[aff1]], %[[aff2]], %[[aff3]], %[[aff4]]] : memref<?x?x?x?xf32>
 //       CHECKPARALLEL:   %[[vb:.*]] = load %[[arg0]][%[[i4]], %[[i5]], %[[i6]], %[[i7]]] : memref<?x?x?x?xf32>
 //       CHECKPARALLEL:   %[[vc:.*]] = load %[[arg2]][%[[i0]], %[[i1]], %[[i2]], %[[i3]]] : memref<?x?x?x?xf32>
 //       CHECKPARALLEL:   %[[inc:.*]] = mulf %[[va]], %[[vb]] : f32
 //       CHECKPARALLEL:   %[[res:.*]] = addf %[[vc]], %[[inc]] : f32
 //       CHECKPARALLEL:   store %[[res]], %[[arg2]][%[[i0]], %[[i1]], %[[i2]], %[[i3]]] : memref<?x?x?x?xf32>