[mlir][sparse] hoist loop invariant tensor loads in sparse compiler

[aartbik]

Summary

After bufferization, the backend has much more trouble hoisting loop invariant
loads from the loops generated by the sparse compiler. Therefore, this is done
during sparse code generation. Note that we don't bother hoisting derived
invariant expressions on SSA values, since the backend does that very well.

Still TBD: scalarize reductions to avoid load-add-store cycles

[aartbik]

rebased

[penpornk]

I'm sorry for the delay! The generated MLIR code looks good to me.

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

68

Nit: It seems inner-level comments should start with // instead of ///. If this is true, please fix. Same for lines 71-72 and 86.

72

Thank you for more code comments!

76

Nit: s/consist/consists/

515

Nit: Why the question mark?

565

Just curious why you refactored this. Do you plan to expand this function in the future?

994

This is clever. :) I was wondering why if (!codegen.loops[idx]) in line 606 would work before I reach this point.

[aartbik]

Thanks for the review!

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

68

For classes/struct, the methods/fields still get the /, inside code it gets
But now that you point this out, I see I missed a few /// cases! Thanks!

515

Ah, because I test that in the next line. Rephrased to be more clear.

565

Yes! The next or after next CL will be vectorization, and for that it is nice to have all the buffer references in a method.

994

Yes, before entering the next loop at the same level, we need to make sure none of the hoisting is valid anymore. So by keeping this "invariant" (pun intended) during the recursion, it all works well together.

[aartbik]

addressed comments

[aartbik]

rebased

[aartbik]

moved to main (I think....)

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

 enum class Kind { kTensor, kInvariant, kMulF, kMulI, kAddF, kAddI };
 
 /// Tensor expression. Represents a MLIR expression in tensor index notation.
 /// For tensors, e0 denotes the tensor index. For invariants, the IR value is
 /// stored directly. For binary operations, e0 and e1 denote the index of the
 /// children tensor expressions.
 struct TensorExp {
   TensorExp(Kind k, unsigned x, unsigned y, Value v)
-      : kind(k), e0(x), e1(y), val(v) {}
+      : kind(k), e0(x), e1(y), val(v) {
+    assert((kind == Kind::kTensor && e0 != -1u && e1 == -1u && !val) ||
+           (kind == Kind::kInvariant && e0 == -1u && e1 == -1u && val) ||
+           (kind >= Kind::kMulF && e0 != -1u && e1 != -1u && !val));
+  }
   Kind kind;
+  /// Indices of children expression(s).
   unsigned e0;
   unsigned e1;
+  /// Direct link to IR for an invariant. During code generation,
+  /// field is used to cache "hoisted" loop invariant tensor loads.
   Value val;
 };
 
 /// Lattice point. Each lattice point consist of a conjunction of tensor
 /// loop indices (encoded in a bitvector) and the index of the corresponding
 /// tensor expression.
 struct LatPoint {
   LatPoint(unsigned n, unsigned e, unsigned b) : bits(n, false), exp(e) {
     bits.set(b);
   }
   LatPoint(const llvm::BitVector &b, unsigned e) : bits(b), exp(e) {}
+  // Conjunction of tensor loop indices as bitvector.
   llvm::BitVector bits;
+  /// Index of the tensor expresssion.
   unsigned exp;
 };
 
 /// A class to handle all iteration lattice operations. This class abstracts
 /// away from some implementation details of storing iteration lattices and
 /// tensor expressions. This allows for fine-tuning performance characteristics
 /// independently from the basic algorithm if bottlenecks are identified.
 class Merger {
if (allDense) {
       codegen.buffers[t] = rewriter.create<AllocaOp>(loc, sparseTp, unknown);
     }
   }
 }
 
 /// Generates a load on a dense or sparse tensor.
 static Value genTensorLoad(Merger &merger, CodeGen &codegen,
                            PatternRewriter &rewriter, linalg::GenericOp op,
-                           unsigned tensor) {
+                           unsigned exp) {
+  Value val = merger.exp(exp).val;
+  // Hoisted invariant load?
+  if (val) {
+    merger.exp(exp).val = Value(); // reset
+    return val;
+  }
+  // Actual load.
   SmallVector<Value, 4> args;
+  unsigned tensor = merger.exp(exp).e0;
   auto map = op.getIndexingMap(tensor);
   bool sparse = false;
   for (unsigned i = 0, m = map.getNumResults(); i < m; ++i) {
     unsigned idx = map.getDimPosition(i);
     args.push_back(codegen.loops[idx]); // universal dense index
     if (sparse || merger.isSparseAccess(tensor, idx)) {
       sparse = true;
       args.clear();
       args.push_back(codegen.pidxs[tensor][idx]); // position index
     }
   }
-  return rewriter.create<LoadOp>(op.getLoc(), codegen.buffers[tensor], args);
+  Location loc = op.getLoc();
+  Value ptr = codegen.buffers[tensor];
+  return rewriter.create<LoadOp>(loc, ptr, args);
 }
 
 /// Generates a store on a dense tensor.
 static void genTensorStore(Merger &merger, CodeGen &codegen,
                            PatternRewriter &rewriter, linalg::GenericOp op,
                            unsigned tensor, Value rhs) {
   SmallVector<Value, 4> args;
   auto map = op.getIndexingMap(tensor);
   for (unsigned i = 0, m = map.getNumResults(); i < m; ++i) {
     unsigned idx = map.getDimPosition(i);
     args.push_back(codegen.loops[idx]); // universal dense index
   }
-  rewriter.create<StoreOp>(op.getLoc(), rhs, codegen.buffers[tensor], args);
+  Location loc = op.getLoc();
+  Value ptr = codegen.buffers[tensor];
+  rewriter.create<StoreOp>(loc, rhs, ptr, args);
 }
 
 /// Generates a pointer/index load from the sparse storage scheme.
-static Value genIntLoad(PatternRewriter &rewriter, Location loc, Value ptr,
-                        Value s) {
+static Value genLoad(PatternRewriter &rewriter, Location loc, Value ptr,
   Value load = rewriter.create<LoadOp>(loc, ptr, s);
   return load.getType().isa<IndexType>()
              ? load
              : rewriter.create<IndexCastOp>(loc, load, rewriter.getIndexType());
 }
 
+/// Generates an invariant value.
+static Value genInvariantValue(Merger &merger, CodeGen &codegen,
+                               PatternRewriter &rewriter, unsigned exp) {
+  return merger.exp(exp).val;
+}
+
 /// Recursively generates tensor expression.
 static Value genExp(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
                     linalg::GenericOp op, unsigned exp) {
   if (merger.exp(exp).kind == Kind::kTensor)
-    return genTensorLoad(merger, codegen, rewriter, op, merger.exp(exp).e0);
+    return genTensorLoad(merger, codegen, rewriter, op, exp);
   else if (merger.exp(exp).kind == Kind::kInvariant)
-    return merger.exp(exp).val;
+    return genInvariantValue(merger, codegen, rewriter, exp);
   Value v0 = genExp(merger, codegen, rewriter, op, merger.exp(exp).e0);
   Value v1 = genExp(merger, codegen, rewriter, op, merger.exp(exp).e1);
   switch (merger.exp(exp).kind) {
   case Kind::kTensor:
   case Kind::kInvariant:
     llvm_unreachable("handled above");
   case Kind::kMulF:
     return rewriter.create<MulFOp>(op.getLoc(), v0, v1);
   case Kind::kMulI:
     return rewriter.create<MulIOp>(op.getLoc(), v0, v1);
   case Kind::kAddF:
     return rewriter.create<AddFOp>(op.getLoc(), v0, v1);
   case Kind::kAddI:
     return rewriter.create<AddIOp>(op.getLoc(), v0, v1);
   }
 }
 
+/// Hoists loop invariant tensor loads for which indices have been exhausted.
+static void genInvariants(Merger &merger, CodeGen &codegen,
+                          PatternRewriter &rewriter, linalg::GenericOp op,
+                          unsigned exp) {
+  if (merger.exp(exp).kind == Kind::kTensor) {
+    unsigned lhs = op.getNumInputsAndOutputs() - 1;
+    unsigned tensor = merger.exp(exp).e0;
+    if (tensor == lhs)
+      return; // TODO: scalarize reduction as well (using scf.yield)
+    auto map = op.getIndexingMap(tensor);
+    for (unsigned i = 0, m = map.getNumResults(); i < m; ++i) {
+      unsigned idx = map.getDimPosition(i);
+      if (!codegen.loops[idx])
+        return; // still in play
+    }
+    // All exhausted at this level.
+    merger.exp(exp).val = genTensorLoad(merger, codegen, rewriter, op, exp);
+
+  } else if (merger.exp(exp).kind != Kind::kInvariant) {
+    // Traverse into the binary operations. Note that we only hoist
+    // tensor loads, since subsequent MLIR/LLVM passes know how to
+    // deal with all other kinds of derived loop invariants.
+    genInvariants(merger, codegen, rewriter, op, merger.exp(exp).e0);
+    genInvariants(merger, codegen, rewriter, op, merger.exp(exp).e1);
+  }
+}
+
 /// Generates initialization code for the subsequent loop sequence at
 /// current index level. Returns true if the loop sequence needs to
 /// maintain the universal index.
 static bool genInit(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
                     linalg::GenericOp op, std::vector<unsigned> &topSort,
                     unsigned at, llvm::BitVector &inits) {
   bool needsUniv = false;
   Location loc = op.getLoc();
if (inits[b]) {
         for (; pat != 0; pat--) {
           if (codegen.pidxs[tensor][topSort[pat - 1]])
             break;
         }
         Value ptr = codegen.pointers[tensor][idx];
         Value one = rewriter.create<ConstantIndexOp>(loc, 1);
         Value p0 = (pat == 0) ? rewriter.create<ConstantIndexOp>(loc, 0)
                               : codegen.pidxs[tensor][topSort[pat - 1]];
-        codegen.pidxs[tensor][idx] = genIntLoad(rewriter, loc, ptr, p0);
+        codegen.pidxs[tensor][idx] = genLoad(rewriter, loc, ptr, p0);
         Value p1 = rewriter.create<AddIOp>(loc, p0, one);
-        codegen.highs[tensor][idx] = genIntLoad(rewriter, loc, ptr, p1);
+        codegen.highs[tensor][idx] = genLoad(rewriter, loc, ptr, p1);
       } else {
         // Dense index still in play.
         needsUniv = true;
       }
     }
   }
 
   // Initialize the universal dense index.
   codegen.loops[idx] = rewriter.create<ConstantIndexOp>(loc, 0);
   return needsUniv;
 }
 
 /// Generates a for-loop on a single index.
 static Operation *genFor(Merger &merger, CodeGen &codegen,
                          PatternRewriter &rewriter, linalg::GenericOp op,
-                         bool isOuter, unsigned idx, llvm::BitVector &indices) {
+                         bool isOuter, bool isInner, unsigned idx,
+                         llvm::BitVector &indices) {
   unsigned fb = indices.find_first();
   unsigned tensor = merger.tensor(fb);
   assert(idx == merger.index(fb));
 
   // Parallelization strategy. Any implicit loop in the Linalg operation that
   // is marked "parallel" is a candidate. Whether it is actually converted to
   // a parallel operation depends on the requested strategy.
   auto iteratorTypes = op.iterator_types().getValue();
static Operation *genWhile(Merger &merger, CodeGen &codegen,
   rewriter.setInsertionPointToStart(&whileOp.after().front());
   return whileOp;
 }
 
 /// Generates a for-loop or a while-loop, depending on whether it implements
 /// singleton iteration or co-iteration over the given conjunction.
 static Operation *genLoop(Merger &merger, CodeGen &codegen,
                           PatternRewriter &rewriter, linalg::GenericOp op,
-                          bool isOuter, unsigned idx, bool needsUniv,
-                          llvm::BitVector &indices) {
-  if (indices.count() == 1)
-    return genFor(merger, codegen, rewriter, op, isOuter, idx, indices);
+                          std::vector<unsigned> &topSort, unsigned at,
+                          bool needsUniv, llvm::BitVector &indices) {
+  unsigned idx = topSort[at];
+  if (indices.count() == 1) {
+    bool isOuter = at == 0;
+    bool isInner = at == topSort.size() - 1;
+    return genFor(merger, codegen, rewriter, op, isOuter, isInner, idx,
+                  indices);
+  }
   return genWhile(merger, codegen, rewriter, op, idx, needsUniv, indices);
 }
 
 /// Generates the local variables for this loop, consisting of the sparse
 /// indices, restored universal dense index, and dense positions.
 static void genLocals(Merger &merger, CodeGen &codegen,
                       PatternRewriter &rewriter, linalg::GenericOp op,
                       std::vector<unsigned> &topSort, unsigned at,
                       bool needsUniv, llvm::BitVector &locals) {
   Location loc = op.getLoc();
   unsigned idx = topSort[at];
 
   // Initialize sparse indices.
   Value min;
   for (unsigned b = 0, be = locals.size(); b < be; b++) {
     if (locals[b] && merger.isSparseBit(b)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value ptr = codegen.indices[tensor][idx];
       Value s = codegen.pidxs[tensor][idx];
-      Value load = genIntLoad(rewriter, loc, ptr, s);
+      Value load = genLoad(rewriter, loc, ptr, s);
       codegen.idxs[tensor][idx] = load;
       if (!needsUniv) {
         if (min) {
           Value cmp =
               rewriter.create<CmpIOp>(loc, CmpIPredicate::ult, load, min);
           min = rewriter.create<SelectOp>(loc, cmp, load, min);
         } else {
           min = load;
static void genStmt(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
   // We maintain the universal dense index if dense indices are still
   // in play for a non-singleton loop sequence.
   unsigned idx = topSort[at];
   unsigned lts = merger.optimize(buildLattices(merger, op, exp, idx));
   unsigned lsize = merger.set(lts).size();
   assert(lsize != 0);
   unsigned l0 = merger.set(lts)[0];
   LatPoint lat0 = merger.lat(l0);
+  genInvariants(merger, codegen, rewriter, op, exp);
   bool needsUniv =
       genInit(merger, codegen, rewriter, op, topSort, at, lat0.bits) &&
       lsize > 1;
 
   // Emit a loop for every lattice point L0 >= Li.
   for (unsigned li : merger.set(lts)) {
     LatPoint lati = merger.lat(li);
 
     // Emit loop.
     llvm::BitVector indices = lati.bits;
     optimizeIndices(merger, lsize, indices);
-    bool isOuter = at == 0;
-    Operation *loop = genLoop(merger, codegen, rewriter, op, isOuter, idx,
-                              needsUniv, indices);
+    Operation *loop =
+        genLoop(merger, codegen, rewriter, op, topSort, at, needsUniv, indices);
     genLocals(merger, codegen, rewriter, op, topSort, at, needsUniv, lati.bits);
 
     // Visit all lattices points with Li >= Lj to generate the
     // loop-body, possibly with if statements for coiteration.
     bool isWhile = dyn_cast<scf::WhileOp>(loop) != nullptr;
     scf::IfOp ifOp;
     for (unsigned lj : merger.set(lts)) {
       if (li == lj || merger.latGT(li, lj)) {
if (isWhile) {
       rewriter.setInsertionPointToEnd(&whileOp.after().front());
       genWhileInduction(merger, codegen, rewriter, op, idx, needsUniv,
                         lati.bits, whileOp.results());
     } else {
       needsUniv = false;
     }
     rewriter.setInsertionPointAfter(loop);
   }
+  codegen.loops[idx] = Value();
 }
 
 namespace {
 
 /// Sparse rewriting rule for generic Lingalg operation.
 struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
 public:
   GenericOpSparsifier(MLIRContext *context, linalg::SparsificationOptions o)

mlir/test/Dialect/Linalg/sparse_2d.mlir

sparse = [
     [ "D","S" ],  // a
     [ ]           // x
   ],
   iterator_types = ["reduction", "reduction"],
   doc = "x = SUM_ij a(i,j)"
 }
 
 // CHECK-LABEL:   func @sum_reduction(
-// CHECK-SAME:                %[[VAL_0:.*]]: tensor<10x20xf32>,
-// CHECK-SAME:                %[[VAL_1:.*]]: tensor<f32>) -> tensor<f32> {
+// CHECK-SAME:                        %[[VAL_0:.*0]]: tensor<10x20xf32>,
+// CHECK-SAME:                        %[[VAL_1:.*1]]: tensor<f32>) -> tensor<f32> {
 // CHECK:           %[[VAL_2:.*]] = constant 999 : index
 // CHECK:           %[[VAL_3:.*]] = constant 10 : index
 // CHECK:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK:           %[[VAL_5:.*]] = constant 1 : index
 // CHECK:           %[[VAL_6:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
 // CHECK:           %[[VAL_7:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
 // CHECK:           %[[VAL_8:.*]] = alloca(%[[VAL_2]]) : memref<?xf32>
 // CHECK:           %[[VAL_9:.*]] = alloca() : memref<f32>

 // CHECK:           %[[VAL_18:.*]] = dim %[[VAL_3]], %[[VAL_5]] : tensor<?x?xf32>
 // CHECK:           %[[VAL_19:.*]] = dim %[[VAL_3]], %[[VAL_6]] : tensor<?x?xf32>
 // CHECK:           %[[VAL_20:.*]] = alloca(%[[VAL_18]], %[[VAL_19]]) : memref<?x?xf32>
 // CHECK:           %[[VAL_21:.*]] = load %[[VAL_7]]{{\[}}%[[VAL_5]]] : memref<?xindex>
 // CHECK:           %[[VAL_22:.*]] = load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
 // CHECK:           scf.for %[[VAL_23:.*]] = %[[VAL_21]] to %[[VAL_22]] step %[[VAL_6]] {
 // CHECK:             %[[VAL_24:.*]] = load %[[VAL_8]]{{\[}}%[[VAL_23]]] : memref<?xindex>
 // CHECK:             scf.for %[[VAL_25:.*]] = %[[VAL_5]] to %[[VAL_15]] step %[[VAL_6]] {
-// CHECK:               %[[VAL_26:.*]] = load %[[VAL_9]]{{\[}}%[[VAL_23]]] : memref<?xindex>
-// CHECK:               %[[VAL_27:.*]] = addi %[[VAL_23]], %[[VAL_6]] : index
-// CHECK:               %[[VAL_28:.*]] = load %[[VAL_9]]{{\[}}%[[VAL_27]]] : memref<?xindex>
-// CHECK:               scf.for %[[VAL_29:.*]] = %[[VAL_26]] to %[[VAL_28]] step %[[VAL_6]] {
-// CHECK:                 %[[VAL_30:.*]] = load %[[VAL_10]]{{\[}}%[[VAL_29]]] : memref<?xindex>
-// CHECK:                 %[[VAL_31:.*]] = load %[[VAL_20]]{{\[}}%[[VAL_24]], %[[VAL_30]]] : memref<?x?xf32>
-// CHECK:                 %[[VAL_32:.*]] = load %[[VAL_11]]{{\[}}%[[VAL_29]]] : memref<?xf32>
-// CHECK:                 %[[VAL_33:.*]] = load %[[VAL_14]]{{\[}}%[[VAL_24]], %[[VAL_25]]] : memref<?x?xf32>
-// CHECK:                 %[[VAL_34:.*]] = load %[[VAL_17]]{{\[}}%[[VAL_25]], %[[VAL_30]]] : memref<?x?xf32>
-// CHECK:                 %[[VAL_35:.*]] = mulf %[[VAL_33]], %[[VAL_34]] : f32
-// CHECK:                 %[[VAL_36:.*]] = mulf %[[VAL_32]], %[[VAL_35]] : f32
-// CHECK:                 %[[VAL_37:.*]] = addf %[[VAL_31]], %[[VAL_36]] : f32
-// CHECK:                 store %[[VAL_37]], %[[VAL_20]]{{\[}}%[[VAL_24]], %[[VAL_30]]] : memref<?x?xf32>
+// CHECK:               %[[VAL_26:.*]] = load %[[VAL_14]]{{\[}}%[[VAL_24]], %[[VAL_25]]] : memref<?x?xf32>
+// CHECK:               %[[VAL_27:.*]] = load %[[VAL_9]]{{\[}}%[[VAL_23]]] : memref<?xindex>
+// CHECK:               %[[VAL_28:.*]] = addi %[[VAL_23]], %[[VAL_6]] : index
+// CHECK:               %[[VAL_29:.*]] = load %[[VAL_9]]{{\[}}%[[VAL_28]]] : memref<?xindex>
+// CHECK:               scf.for %[[VAL_30:.*]] = %[[VAL_27]] to %[[VAL_29]] step %[[VAL_6]] {
+// CHECK:                 %[[VAL_31:.*]] = load %[[VAL_10]]{{\[}}%[[VAL_30]]] : memref<?xindex>
+// CHECK:                 %[[VAL_32:.*]] = load %[[VAL_20]]{{\[}}%[[VAL_24]], %[[VAL_31]]] : memref<?x?xf32>
+// CHECK:                 %[[VAL_33:.*]] = load %[[VAL_11]]{{\[}}%[[VAL_30]]] : memref<?xf32>
+// CHECK:                 %[[VAL_34:.*]] = load %[[VAL_17]]{{\[}}%[[VAL_25]], %[[VAL_31]]] : memref<?x?xf32>
+// CHECK:                 %[[VAL_35:.*]] = mulf %[[VAL_26]], %[[VAL_34]] : f32
+// CHECK:                 %[[VAL_36:.*]] = mulf %[[VAL_33]], %[[VAL_35]] : f32
+// CHECK:                 %[[VAL_37:.*]] = addf %[[VAL_32]], %[[VAL_36]] : f32
+// CHECK:                 store %[[VAL_37]], %[[VAL_20]]{{\[}}%[[VAL_24]], %[[VAL_31]]] : memref<?x?xf32>
 // CHECK:               }
 // CHECK:             }
 // CHECK:           }
 // CHECK:           %[[VAL_38:.*]] = tensor_load %[[VAL_20]] : memref<?x?xf32>
 // CHECK:           return %[[VAL_38]] : tensor<?x?xf32>
 // CHECK:         }
 func @sampled_dense_dense(%args: tensor<?x?xf32>,
                           %arga: tensor<?x?xf32>,

mlir/test/Dialect/Linalg/sparse_3d.mlir

 // CHECK:             scf.for %[[VAL_20:.*]] = %[[VAL_5]] to %[[VAL_10]] step %[[VAL_6]] {
 // CHECK:               %[[VAL_21:.*]] = muli %[[VAL_10]], %[[VAL_19]] : index
 // CHECK:               %[[VAL_22:.*]] = addi %[[VAL_21]], %[[VAL_20]] : index
 // CHECK:               %[[VAL_23:.*]] = load %[[VAL_7]]{{\[}}%[[VAL_22]]] : memref<?xindex>
 // CHECK:               %[[VAL_24:.*]] = addi %[[VAL_22]], %[[VAL_6]] : index
 // CHECK:               %[[VAL_25:.*]] = load %[[VAL_7]]{{\[}}%[[VAL_24]]] : memref<?xindex>
 // CHECK:               scf.for %[[VAL_26:.*]] = %[[VAL_23]] to %[[VAL_25]] step %[[VAL_6]] {
 // CHECK:                 %[[VAL_27:.*]] = load %[[VAL_8]]{{\[}}%[[VAL_26]]] : memref<?xindex>
-// CHECK:                 scf.for %[[VAL_28:.*]] = %[[VAL_5]] to %[[VAL_17]] step %[[VAL_6]] {
-// CHECK:                   %[[VAL_29:.*]] = load %[[VAL_9]]{{\[}}%[[VAL_26]]] : memref<?xf32>
-// CHECK:                   %[[VAL_30:.*]] = load %[[VAL_12]]{{\[}}%[[VAL_20]], %[[VAL_28]]] : memref<?x?xf32>
-// CHECK:                   %[[VAL_31:.*]] = mulf %[[VAL_29]], %[[VAL_30]] : f32
-// CHECK:                   %[[VAL_32:.*]] = load %[[VAL_15]]{{\[}}%[[VAL_27]], %[[VAL_28]]] : memref<?x?xf32>
+// CHECK:                 %[[VAL_28:.*]] = load %[[VAL_9]]{{\[}}%[[VAL_26]]] : memref<?xf32>
+// CHECK:                 scf.for %[[VAL_29:.*]] = %[[VAL_5]] to %[[VAL_17]] step %[[VAL_6]] {
+// CHECK:                   %[[VAL_30:.*]] = load %[[VAL_12]]{{\[}}%[[VAL_20]], %[[VAL_29]]] : memref<?x?xf32>
+// CHECK:                   %[[VAL_31:.*]] = mulf %[[VAL_28]], %[[VAL_30]] : f32
+// CHECK:                   %[[VAL_32:.*]] = load %[[VAL_15]]{{\[}}%[[VAL_27]], %[[VAL_29]]] : memref<?x?xf32>
 // CHECK:                   %[[VAL_33:.*]] = mulf %[[VAL_31]], %[[VAL_32]] : f32
-// CHECK:                   %[[VAL_34:.*]] = load %[[VAL_18]]{{\[}}%[[VAL_19]], %[[VAL_28]]] : memref<?x?xf32>
+// CHECK:                   %[[VAL_34:.*]] = load %[[VAL_18]]{{\[}}%[[VAL_19]], %[[VAL_29]]] : memref<?x?xf32>
 // CHECK:                   %[[VAL_35:.*]] = addf %[[VAL_33]], %[[VAL_34]] : f32
-// CHECK:                   store %[[VAL_35]], %[[VAL_18]]{{\[}}%[[VAL_19]], %[[VAL_28]]] : memref<?x?xf32>
+// CHECK:                   store %[[VAL_35]], %[[VAL_18]]{{\[}}%[[VAL_19]], %[[VAL_29]]] : memref<?x?xf32>
 // CHECK:                 }
 // CHECK:               }
 // CHECK:             }
 // CHECK:           }
 // CHECK:           %[[VAL_36:.*]] = tensor_load %[[VAL_18]] : memref<?x?xf32>
 // CHECK:           return %[[VAL_36]] : tensor<?x?xf32>
 // CHECK:         }
 func @kernel_3d(%arga: tensor<?x?xf32>,
%0 = linalg.generic #trait_sum_reduction
     ins(%arga : tensor<10x20x30xf32>)
     init(%argx : tensor<f32>) {
       ^bb(%a : f32, %x : f32):
         %0 = addf %x, %a  : f32
         linalg.yield %0: f32
   } -> tensor<f32>
   return %0 : tensor<f32>
 }
+
+#trait_invariants = {
+  indexing_maps = [
+    affine_map<(i,j,k) -> (i)>,      // a
+    affine_map<(i,j,k) -> (j)>,      // b
+    affine_map<(i,j,k) -> (k)>,      // c
+    affine_map<(i,j,k) -> (i,j,k)>   // x
+  ],
+  sparse = [
+    [ "D" ],           // a
+    [ "D" ],           // b
+    [ "D" ],           // c
+    [ "D", "D", "D" ]  // x
+  ],
+  iterator_types = ["parallel", "parallel", "parallel"],
+  doc = "x(i,j,k) = a(i) * b(j) * c(k)"
+}
+
+// CHECK-LABEL:   func @invariants(
+// CHECK-SAME:                     %[[VAL_0:.*]]: tensor<10xf32>,
+// CHECK-SAME:                     %[[VAL_1:.*]]: tensor<20xf32>,
+// CHECK-SAME:                     %[[VAL_2:.*]]: tensor<30xf32>) -> tensor<10x20x30xf32> {
+// CHECK:           %[[VAL_3:.*]] = constant 10 : index
+// CHECK:           %[[VAL_4:.*]] = constant 20 : index
+// CHECK:           %[[VAL_5:.*]] = constant 30 : index
+// CHECK:           %[[VAL_6:.*]] = constant 0 : index
+// CHECK:           %[[VAL_7:.*]] = constant 1 : index
+// CHECK:           %[[VAL_8:.*]] = alloca() : memref<10xf32>
+// CHECK:           %[[VAL_9:.*]] = alloca() : memref<20xf32>
+// CHECK:           %[[VAL_10:.*]] = alloca() : memref<30xf32>
+// CHECK:           %[[VAL_11:.*]] = alloca() : memref<10x20x30xf32>
+// CHECK:           scf.for %[[VAL_12:.*]] = %[[VAL_6]] to %[[VAL_3]] step %[[VAL_7]] {
+// CHECK:             %[[VAL_13:.*]] = load %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<10xf32>
+// CHECK:             scf.for %[[VAL_14:.*]] = %[[VAL_6]] to %[[VAL_4]] step %[[VAL_7]] {
+// CHECK:               %[[VAL_15:.*]] = load %[[VAL_9]]{{\[}}%[[VAL_14]]] : memref<20xf32>
+// CHECK:               scf.for %[[VAL_16:.*]] = %[[VAL_6]] to %[[VAL_5]] step %[[VAL_7]] {
+// CHECK:                 %[[VAL_17:.*]] = mulf %[[VAL_13]], %[[VAL_15]] : f32
+// CHECK:                 %[[VAL_18:.*]] = load %[[VAL_10]]{{\[}}%[[VAL_16]]] : memref<30xf32>
+// CHECK:                 %[[VAL_19:.*]] = mulf %[[VAL_17]], %[[VAL_18]] : f32
+// CHECK:                 store %[[VAL_19]], %[[VAL_11]]{{\[}}%[[VAL_12]], %[[VAL_14]], %[[VAL_16]]] : memref<10x20x30xf32>
+// CHECK:               }
+// CHECK:             }
+// CHECK:           }
+// CHECK:           %[[VAL_20:.*]] = tensor_load %[[VAL_11]] : memref<10x20x30xf32>
+// CHECK:           return %[[VAL_20]] : tensor<10x20x30xf32>
+// CHECK:         }
+func @invariants(%arga: tensor<10xf32>,
+                 %argb: tensor<20xf32>,
+                 %argc: tensor<30xf32>) -> tensor<10x20x30xf32> {
+  %0 = linalg.generic #trait_invariants
+    ins(%arga, %argb, %argc : tensor<10xf32>, tensor<20xf32>, tensor<30xf32>) {
+      ^bb(%a : f32, %b : f32, %c : f32):
+        %0 = mulf %a, %b  : f32
+        %1 = mulf %0, %c  : f32
+        linalg.yield %1: f32
+  } -> tensor<10x20x30xf32>
+  return %0 : tensor<10x20x30xf32>
+}