[mlir][Linalg] Revisit 0-D abstraction

[nicolasvasilache]

Summary

This revision takes advantage of the empty AffineMap to specify the
0-D edge case. This allows removing a bunch of annoying corner cases
that ended up impacting users of Linalg.

[rriddle]

mlir/lib/IR/AffineMap.cpp

328

Do we need these asserts? None of the other methods assert this.

353

Same here.

[nicolasvasilache]

Remove unnecessary asserts.

[pifon2a]

Awesome! Thank you for doing this, @nicolasvasilache!

[mravishankar]

This is awesome, but this is a breaking change to clients who had implemented the WAR. I can help fix the IREE changes.

[herhut]

Thanks, much cleaner solution!

[nicolasvasilache]

ok since everyone seems to be onboard and pifon@ is integrate duty this week I'll coordinate with him.
Thanks for your quick turnaround!

[herhut]

In D75831#1912206, @nicolasvasilache wrote:

ok since everyone seems to be onboard and pifon@ is integrate duty this week I'll coordinate with him.

pifon@ is next week.

[mravishankar]

I looked at this change a little bit deeper, and I am not sure this works. This seems to have an affine_map for scalars being () -> (). This seems odd. Isn't it better to have a scalar access be (i,j) -> (), or something like that.
That needs change to the AffineMap constructor above.

[nicolasvasilache]

This seems to have an affine_map for scalars being () -> (). This seems odd. Isn't it better to have a scalar access be (i,j) -> (), or something like that.

Whatever we do for 0-D is a convention, we can chose the convention we want.
()->() did not require parser changes whereas (i, j) -> () does.

I went ahead and implemented those changes.

[nicolasvasilache]

Address comments.

[bondhugula]

mlir/include/mlir/IR/AffineMap.h

49 ↗

(On Diff #249358)

We could just go for for:

static AffineMap get(unsigned dimCount, unsigned symCount, MLIRContext *context);

In any case, please have context appear last to be consistent with everything else.

[bondhugula]

mlir/include/mlir/IR/AffineMap.h

47 ↗

(On Diff #249358)

dim -> dimensions

mlir/lib/IR/AffineMap.cpp

367

/*numSymbols=*/0

367

How do you know whether results is non-empty here?

[nicolasvasilache]

Address review.

[nicolasvasilache]

Address comments.

[nicolasvasilache]

mlir/lib/IR/AffineMap.cpp

367

Good point, thanks Uday!

[aartbik]

Very cool to see this happen! This will allow me to pay off some engineering debt for vector.contract as well (where I was forced to use empty maps rather than (i,j) -> () maps).

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

llvm::Optional<SmallVector<StringRef, 8>> referenceIterators() {
       return SmallVector<StringRef, 8>{getReductionIteratorTypeName()};
     }
 
     // A(r_i) * B(r_i) -> C()
     llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
       MLIRContext *context = getContext();
       auto r_i = getAffineDimExpr(0, context);
       return SmallVector<AffineMap, 8>{
-        AffineMap::get(1, 0, {r_i}), AffineMap::get(1, 0, {r_i}), AffineMap()};
+        AffineMap::get(1, 0, {r_i}),
+        AffineMap::get(1, 0, {r_i}),
+        AffineMap::get(context)};
     }
   }];
 
   let hasFolder = 1;
 }
 
 def MatvecOp : LinalgStructured_Op<"matvec", [NInputs<2>, NOutputs<1>]> {
 

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

for (auto en : llvm::enumerate(op.indexing_maps())) {
     indexingMaps.push_back(m); // Save reference to map for further checks.
     auto view = (idx < nInputViews) ? op.getInputShapedType(idx)
                                     : op.getOutputShapedType(idx - nInputViews);
 
     if (m.getNumSymbols() != 0)
       return op.emitOpError("expected indexing_map #")
              << idx << " to have no symbols";
 
-    if (m.getNumDims() != nLoops)
+    if (!m.isEmpty() && m.getNumDims() != nLoops)
       return op.emitOpError("expected indexing_map #")
              << idx << " to have " << nLoops
              << " dim(s) to match the number of loops";
 
-    if (m.getNumResults() == 1 && view.getRank() == 0) {
-      auto cst = m.getResult(0).template dyn_cast<AffineConstantExpr>();
-      if (!cst || cst.getValue() != 0)
-        return op.emitOpError("expected indexing_map #")
-               << idx << " to be 0 to match 0-D view: " << view;
-    } else if (m.getNumResults() != view.getRank()) {
+    if (m.getNumResults() != view.getRank())
       return op.emitOpError("expected indexing_map #")
              << idx << " results to match view rank: " << view;
-    }
-  }
 
   auto concatMap = concatAffineMaps(indexingMaps);
   auto aggregateMap = inversePermutation(concatMap);
   if (!aggregateMap)
     return op.emitOpError("expected the concatenation of maps in indexing_map "
                           "to be invertible");
 
   return success();

 } // namespace mlir
 
 AffineMap mlir::linalg::extractOrIdentityMap(Optional<AffineMap> maybeMap,
                                              unsigned rank,
                                              MLIRContext *context) {
   if (maybeMap)
     return maybeMap.getValue();
   if (rank == 0)
-    return AffineMap();
+    return AffineMap::get(context);
   return AffineMap::getMultiDimIdentityMap(rank, context);
 }
 
 SmallVector<AffineExpr, 4>
 mlir::linalg::makeAffineDimExprs(unsigned num, unsigned &startIdx,
                                  MLIRContext *context) {
   SmallVector<AffineExpr, 4> res;
   res.reserve(num);

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

 using namespace mlir::linalg;
 
 using edsc::op::operator+;
 using edsc::op::operator==;
 
 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<ValueHandle, 8> res;
   res.reserve(map.getNumResults());
   auto dims = map.getNumDims();
   for (auto e : map.getResults()) {
     auto exprMap = AffineMap::get(dims, 0, e);
     SmallVector<Value, 4> operands(vals.begin(), vals.end());
static void emitScalarImplementation(ArrayRef<Value> allIvs,
     auto loc = ScopedContext::getLocation();
     using edsc::intrinsics::detail::ValueHandleArray;
     unsigned nInputs = genericOp.getNumInputs();
     unsigned nOutputs = genericOp.getNumOutputs();
     SmallVector<Value, 4> indexedValues(nInputs + nOutputs);
 
     // 1.a. Emit std_load from input views.
     for (unsigned i = 0; i < nInputs; ++i) {
-      Value input = genericOp.getInput(i);
-      if (input.getType().cast<ShapedType>().getRank()) {
-        ValueHandleArray indexing(makeCanonicalAffineApplies(
-            b, loc, genericOp.getInputIndexingMap(i), allIvs));
-        indexedValues[i] = std_load(input, indexing);
-      } else {
-        indexedValues[i] = std_load(input);
-      }
+      indexedValues[i] = std_load(genericOp.getInput(i), indexing);
     }
 
     // 1.b. Emit std_load from output views.
     for (unsigned i = 0; i < nOutputs; ++i) {
       Value output = genericOp.getOutputBuffer(i);
-      if (output.getType().cast<ShapedType>().getRank()) {
-        ValueHandleArray indexing(makeCanonicalAffineApplies(
-            b, loc, genericOp.getOutputIndexingMap(i), allIvs));
-        indexedValues[nInputs + i] = std_load(output, indexing);
-      } else {
-        indexedValues[nInputs + i] = std_load(output);
-      }
     }
 
     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);
-        if (output.getType().cast<ShapedType>().getRank()) {
-          ValueHandleArray indexing(makeCanonicalAffineApplies(
-              b, loc, genericOp.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 = 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);
     for (unsigned i = 0; i < nOutputs; ++i) {
-      Value output = genericOp.getOutputBuffer(i);
-      if (output.getType().cast<ShapedType>().getRank()) {
-        ValueHandleArray indexing(makeCanonicalAffineApplies(
-            b, loc, genericOp.getOutputIndexingMap(i), allIvs));
-        std_store(map.lookup(yieldOp->getOperand(i)),
-                  genericOp.getOutputBuffer(i), indexing);
-      } else {
-        std_store(map.lookup(yieldOp->getOperand(i)), output);
-      }
     }
   }
 };
 
 // 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.

mlir/lib/IR/AffineMap.cpp

AffineMap mlir::simplifyAffineMap(AffineMap map) {
   for (auto e : map.getResults()) {
     exprs.push_back(
         simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols()));
   }
   return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs);
 }
 
 AffineMap mlir::inversePermutation(AffineMap map) {
-  if (!map)
+  assert(map);
+  if (map.isEmpty())
     return map;
   assert(map.getNumSymbols() == 0 && "expected map without symbols");
   SmallVector<AffineExpr, 4> exprs(map.getNumDims());
   for (auto en : llvm::enumerate(map.getResults())) {
     auto expr = en.value();
     // Skip non-permutations.
     if (auto d = expr.dyn_cast<AffineDimExpr>()) {
       if (exprs[d.getPosition()])
         continue;
       exprs[d.getPosition()] = getAffineDimExpr(en.index(), d.getContext());
     }
   }
   SmallVector<AffineExpr, 4> seenExprs;
   seenExprs.reserve(map.getNumDims());
   for (auto expr : exprs)
     if (expr)
       seenExprs.push_back(expr);
   if (seenExprs.size() != map.getNumInputs())
     return AffineMap();
   return AffineMap::get(map.getNumResults(), 0, seenExprs);
 }
 
 AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
+  assert(!maps.empty());
+  assert(llvm::all_of(maps, [](AffineMap m) { return m; }));
   unsigned numResults = 0;
   for (auto m : maps)
-    numResults += (m && !m.isSingleConstant()) ? m.getNumResults() : 0;
+    numResults += m.getNumResults();
   unsigned numDims = 0;
   SmallVector<AffineExpr, 8> results;
   results.reserve(numResults);
   for (auto m : maps) {
-    if (!m || m.isSingleConstant())
-      continue;
     assert(m.getNumSymbols() == 0 && "expected map without symbols");
     results.append(m.getResults().begin(), m.getResults().end());
     numDims = std::max(m.getNumDims(), numDims);
   }
-  return numDims == 0 ? AffineMap() : AffineMap::get(numDims, 0, results);
+  return numDims == 0 ? AffineMap::get(maps.front().getContext())
+                      : AffineMap::get(numDims, 0, results);
 }
 
 //===----------------------------------------------------------------------===//
 // MutableAffineMap.
 //===----------------------------------------------------------------------===//
 
 MutableAffineMap::MutableAffineMap(AffineMap map)
     : numDims(map.getNumDims()), numSymbols(map.getNumSymbols()),

mlir/test/Dialect/Linalg/invalid.mlir

linalg.generic {
     iterator_types = ["parallel"]
   } %arg0: memref<i32>
 }
 
 // -----
 
 func @foo(%0: i32) -> i32 { return %0: i32 }
 
-func @generic_wrong_dim_in_map(%arg0: memref<i32>) {
+func @generic_wrong_dim_in_map(%arg0: memref<1xi32>) {
   // expected-error @+1 {{op expected indexing_map #0 to have 1 dim(s) to match the number of loops}}
   linalg.generic {
     args_in = 0,
     args_out = 1,
     fun = @foo,
     indexing_maps =  [ affine_map<() -> (0)> ],
     iterator_types = ["parallel"]
-  } %arg0: memref<i32>
-}
-
-// -----
-
-func @foo(%0: i32) -> i32 { return %0: i32 }
-
-func @generic_zero_d_view(%arg0: memref<i32>) {
-  // expected-error @+1 {{op expected indexing_map #0 to be 0 to match 0-D view: 'memref<i32>'}}
-  linalg.generic {
-    args_in = 0,
-    args_out = 1,
-    fun = @foo,
-    indexing_maps =  [ affine_map<() -> (1)> ],
-    iterator_types = []
-  } %arg0: memref<i32>
+  } %arg0: memref<1xi32>
 }
 
 // -----
 
 func @foo(%0: f32) -> f32 { return %0: f32 }
 
 func @generic_one_d_view(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
   // expected-error @+1 {{op expected indexing_map #0 results to match view rank: 'memref<?xf32, affine_map<(d0)[s0] -> (d0 + s0)>>'}}

mlir/test/Dialect/Linalg/loops.mlir

 // CHECK:       %[[ijk_float:.*]] = sitofp %[[ijk_int]] : i32 to f32
 // CHECK:       %[[result_2:.*]] = addf %[[c]], %[[ijk_float]] : f32
 // CHECK:       store %[[result_1]], %{{.*}}[%[[i]], %[[j]], %[[k]]]
 // CHECK:       store %[[result_2]], %{{.*}}[%[[i]], %[[k]], %[[j]]]
 
 // -----
 
 #broadcast_access = [
-  affine_map<(i, j) -> (0)>,
+  affine_map<() -> ()>,
   affine_map<(i, j) -> (i, j)>
 ]
 
 #trait_broadcast = {
   args_in = 1,
   args_out = 1,
   indexing_maps = #broadcast_access,
   iterator_types = ["parallel", "parallel"],

 // CHECK:     %[[a:.*]] = load %[[ARG0]][
 // CHECK:     %[[ij:.*]] = addi %[[i]], %[[j]] : index
 // CHECK:     %[[ij_int:.*]] = index_cast %[[ij]] : index to i32
 // CHECK:     %[[result:.*]] = addi %[[a]], %[[ij_int]] : i32
 // CHECK:     store %[[result]], %[[ARG1]][%[[i]], %[[j]]]
 
 #reduce_1D_access = [
   affine_map<(i) -> (i)>,
-  affine_map<(i) -> (0)>
+  affine_map<() -> ()>
 ]
 
 #trait_reduce_1D = {
   args_in = 1,
   args_out = 1,
   indexing_maps = #reduce_1D_access,
   iterator_types = ["reduction"],
   library_call = "some_reduce_external_fn"

 // CHECK:   %[[a:.*]] = load %[[ARG0]][%[[i]]]
 // CHECK:   %[[b:.*]] = load %[[ARG1]][]
 // CHECK:   %[[c:.*]] = addf %[[a]], %[[b]] : f32
 // CHECK:   store %[[c]], %[[ARG1]][]
 
 
 #reduce_init_1D_access = [
   affine_map<(i) -> (i)>,
-  affine_map<(i) -> (0)>,
-  affine_map<(i) -> (0)>
+  affine_map<() -> ()>,
+  affine_map<() -> ()>
 ]
 
 #trait_reduce_init_1D = {
   args_in = 2,
   args_out = 1,
   indexing_maps = #reduce_init_1D_access,
   iterator_types = ["reduction"],
   library_call = "some_reduce_external_fn"

mlir/test/Dialect/Linalg/roundtrip.mlir

 //  CHECK-SAME:     indexing_maps = [#{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel"],
 //  CHECK-SAME:     library_call = "some_external_function_name_1"} %{{.*}}, %{{.*}} {foo = 1 : i64}:
 //  CHECK-SAME:     tensor<?x?xvector<3x4xi4>>, tensor<?x?x?xf32> -> tensor<?x?x?xf32>
 //       CHECK:   return {{.*}} : tensor<?x?x?xf32>
 
 // -----
 
 #broadcast_access = [
-  affine_map<(i, j) -> (0)>,
+  affine_map<() -> ()>,
   affine_map<(i, j) -> (i, j)>
 ]
 
 #trait_broadcast = {
   args_in = 1,
   args_out = 1,
   indexing_maps = #broadcast_access,
   iterator_types = ["parallel", "parallel"],