[mlir][Linalg] NFC: Clean up for 0-D abstraction.

[hanchung]

Summary

After D75831 has been landed, both the generic op and indexed_generic op can
handle 0-D edge case. In the previous patch, only generic op has been updated.
This patch updates the lowering to loops for indexed_generic op. Since they are
almost the sanme, the patch also refactors the common part.

[mravishankar]

Just one suggestion since you are touching these anyways

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

327–329

This is unrelated to the change, but you can avoid doing these loads here if the corresponding argument of the region has no uses. Would be a bonus to do this here.
I think doing that in a separate change is fine. Can you just add a todo for me here if you dont get to it

[rriddle]

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

90

Isn't this returning a reference?

95

Can you add a message to this assert?

100

This looks weird, why not just assert the back is a yield?

101

nit: Please cache the end iterator of the loop. This is the LLVM style.

[hanchung]

Address comments.

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

90

Yes, it is. Thanks!

[hanchung]

Fix build.

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

 
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 using namespace mlir::linalg;
 
 using edsc::op::operator+;
 using edsc::op::operator==;
+using mlir::edsc::intrinsics::detail::ValueHandleArray;
 
 static SmallVector<ValueHandle, 8>
 makeCanonicalAffineApplies(OpBuilder &b, Location loc, AffineMap map,
                            ArrayRef<Value> vals) {
   if (map.isEmpty())
     return {};
   assert(map.getNumSymbols() == 0);
   assert(map.getNumInputs() == vals.size());
SmallVector<Value, 4> emitLoopRanges(OpBuilder &b, Location loc, AffineMap map,
   SmallVector<Value, 4> res;
   for (unsigned idx = 0, e = map.getNumResults(); idx < e; ++idx) {
     res.push_back(
         linalg_range(std_constant_index(0), sizes[idx], std_constant_index(1)));
   }
   return res;
 }
 
+template <typename OpType>
+static void inlineRegionAndEmitStdStore(OpType op,
+                                        ArrayRef<Value> indexedValues,
+                                        ArrayRef<ValueHandleArray> indexing,
+                                        ArrayRef<Value> outputBuffers) {
+  auto &b = ScopedContext::getBuilder();
+  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 an yield op in the end of the region");
+  for (unsigned i = 0, e = terminator.getNumOperands(); i < e; ++i) {
+    std_store(map.lookup(terminator.getOperand(i)), outputBuffers[i],
+              indexing[i]);
+  }
+}
+
 namespace {
 template <typename IndexedValueType, typename LinalgOpType>
 class LinalgScopedEmitter {};
 
 template <typename IndexedValueType>
 class LinalgScopedEmitter<IndexedValueType, CopyOp> {
 public:
   static void emitScalarImplementation(ArrayRef<Value> allIvs, CopyOp copyOp) {
static void emitScalarImplementation(ArrayRef<Value> allIvs,
 
     // 1.a. Emit std_load from input views.
     for (unsigned i = 0; i < nInputs; ++i) {
       ValueHandleArray indexing(makeCanonicalAffineApplies(
           b, loc, genericOp.getInputIndexingMap(i), allIvs));
       indexedValues[i] = std_load(genericOp.getInput(i), indexing);
     }
 
     // 1.b. Emit std_load from output views.
+    // TODO(mravishankar): Avoid the loads if the corresponding argument of the
+    // region has no uses.
     for (unsigned i = 0; i < nOutputs; ++i) {
       Value output = genericOp.getOutputBuffer(i);
       ValueHandleArray indexing(makeCanonicalAffineApplies(
           b, loc, genericOp.getOutputIndexingMap(i), allIvs));
       indexedValues[nInputs + i] = std_load(output, indexing);
     }
 
     auto funcOp = genericOp.getFunction();
     if (funcOp) {
       // 2. Emit call.
       Operation *callOp = std_call(funcOp, indexedValues);
       assert(callOp->getNumResults() == genericOp.getNumOutputs());
 
       // 3. Emit std_store.
       for (unsigned i = 0; i < nOutputs; ++i) {
         Value output = genericOp.getOutputBuffer(i);
         ValueHandleArray indexing(makeCanonicalAffineApplies(
             b, loc, genericOp.getOutputIndexingMap(i), allIvs));
         std_store(callOp->getResult(i), output, indexing);
       }
       return;
     }
     // TODO(ntv): When a region inliner exists, use it.
     // 2. Inline region, currently only works for a single basic block.
-    BlockAndValueMapping map;
-    auto &block = genericOp.region().front();
-    map.map(block.getArguments(), indexedValues);
-    for (auto &op : block.without_terminator()) {
-      assert(op.getNumRegions() == 0);
-      auto *newOp = b.clone(op, map);
-      map.map(op.getResults(), newOp->getResults());
-    }
-
     // 3. Emit std_store.
-    auto *yieldOp = cast<YieldOp>(block.back()).getOperation();
-    assert(yieldOp->getNumOperands() == nOutputs);
+    SmallVector<ValueHandleArray, 8> indexing;
+    SmallVector<Value, 8> outputBuffers;
     for (unsigned i = 0; i < nOutputs; ++i) {
-      ValueHandleArray indexing(makeCanonicalAffineApplies(
+      indexing.emplace_back(makeCanonicalAffineApplies(
           b, loc, genericOp.getOutputIndexingMap(i), allIvs));
-      std_store(map.lookup(yieldOp->getOperand(i)),
-                genericOp.getOutputBuffer(i), indexing);
+      outputBuffers.push_back(genericOp.getOutputBuffer(i));
     }
+    inlineRegionAndEmitStdStore(genericOp, indexedValues, indexing,
+                                outputBuffers);
   }
 };
 
 // Emits the MLIR for the scalar part of the indexed generic op by:
 //   1. Emitting std_load and std_store ops for each input and output view in
 //      order. This is achieved by applying the appropriate input or output map
 //      to the enclosing induction variables.
 //   2. Emitting a call to `op.fun()` that takes as arguments the induction
static void emitScalarImplementation(ArrayRef<Value> allIvs,
 
     for (unsigned i = 0; i < nLoops; ++i) {
       indexedValues[i] = allIvs[i];
     }
 
     // 1.a. Emit std_load from input views.
     for (unsigned i = 0; i < nInputs; ++i) {
       Value input = indexedGenericOp.getInput(i);
-      if (input.getType().cast<ShapedType>().getRank()) {
-        ValueHandleArray indexing(makeCanonicalAffineApplies(
-            b, loc, indexedGenericOp.getInputIndexingMap(i), allIvs));
-        indexedValues[nLoops + i] = std_load(input, indexing);
-      } else {
-        indexedValues[nLoops + i] = std_load(input);
-      }
     }
 
     // 1.b. Emit std_load from output views.
     for (unsigned i = 0; i < nOutputs; ++i) {
       Value output = indexedGenericOp.getOutputBuffer(i);
-      if (output.getType().cast<ShapedType>().getRank()) {
-        ValueHandleArray indexing(makeCanonicalAffineApplies(
-            b, loc, indexedGenericOp.getOutputIndexingMap(i), allIvs));
-        indexedValues[nLoops + nInputs + i] = std_load(output, indexing);
-      } else {
-        indexedValues[nLoops + nInputs + i] = std_load(output);
-      }
     }
 
     if (auto funcOp = indexedGenericOp.getFunction()) {
       // 2. Emit call.
       Operation *callOp = std_call(funcOp, indexedValues);
       assert(callOp->getNumResults() == indexedGenericOp.getNumOutputs());
 
       // 3. Emit std_store.
       for (unsigned i = 0; i < nOutputs; ++i) {
         Value output = indexedGenericOp.getOutputBuffer(i);
-        if (output.getType().cast<ShapedType>().getRank()) {
-          ValueHandleArray indexing(makeCanonicalAffineApplies(
-              b, loc, indexedGenericOp.getOutputIndexingMap(i), allIvs));
-          std_store(callOp->getResult(i), output, indexing);
-        } else {
-          std_store(callOp->getResult(i), output);
-        }
       }
       return;
     }
     // TODO(ntv): When a region inliner exists, use it.
     // 2. Inline region, currently only works for a single basic block.
-    BlockAndValueMapping map;
-    auto &block = indexedGenericOp.region().front();
-    map.map(block.getArguments(), indexedValues);
-    for (auto &op : block.without_terminator()) {
-      assert(op.getNumRegions() == 0);
-      auto *newOp = b.clone(op, map);
-      map.map(op.getResults(), newOp->getResults());
-    }
-
     // 3. Emit std_store.
-    auto *yieldOp = cast<YieldOp>(block.back()).getOperation();
-    assert(yieldOp->getNumOperands() == nOutputs);
+    SmallVector<ValueHandleArray, 8> indexing;
+    SmallVector<Value, 8> outputBuffers;
     for (unsigned i = 0; i < nOutputs; ++i) {
-      Value output = indexedGenericOp.getOutputBuffer(i);
-      if (output.getType().cast<ShapedType>().getRank()) {
-        ValueHandleArray indexing(makeCanonicalAffineApplies(
-            b, loc, indexedGenericOp.getOutputIndexingMap(i), allIvs));
-        std_store(map.lookup(yieldOp->getOperand(i)), output, indexing);
-      } else {
-        std_store(map.lookup(yieldOp->getOperand(i)), output);
-      }
+      indexing.emplace_back(makeCanonicalAffineApplies(
+      outputBuffers.push_back(indexedGenericOp.getOutputBuffer(i));
     }
+    inlineRegionAndEmitStdStore(indexedGenericOp, indexedValues, indexing,
+                                outputBuffers);
   }
 };
 
 // This struct is for factoring out the implementation and support template
 // instantiations in the following 2 cases:
 //   1. Appending to a list of patterns via RewritePatternList.
 //   2. Direct invocation via `linalgOpToLoops` and `linalgOpToAffineLoops`.
 // The implementation must work both in DRR and inside a RewritePattern. As a