[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

329

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

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

dim -> dimensions

mlir/lib/IR/AffineMap.cpp

365

/*numSymbols=*/0

365

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

[nicolasvasilache]

Address review.

[nicolasvasilache]

Address comments.

[nicolasvasilache]

mlir/lib/IR/AffineMap.cpp

365

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/docs/Dialects/Affine.md

affine-expr ::= `(` affine-expr `)`
               | `-`? integer-literal `*` affine-expr
               | affine-expr `ceildiv` integer-literal
               | affine-expr `floordiv` integer-literal
               | affine-expr `mod` integer-literal
               | `-`affine-expr
               | bare-id
               | `-`? integer-literal
 
-multi-dim-affine-expr ::= `(` affine-expr (`,` affine-expr)* `)`
+multi-dim-affine-expr ::= `(` `)`
+                        | `(` affine-expr (`,` affine-expr)* `)`
 ```
 
 `ceildiv` is the ceiling function which maps the result of the division of its
 first argument by its second argument to the smallest integer greater than or
 equal to that result. `floordiv` is a function which maps the result of the
 division of its first argument by its second argument to the largest integer
 less than or equal to that result. `mod` is the modulo operation: since its
 second argument is always positive, its results are always positive in our

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(1, 0, context)};
     }
   }];
 
   let hasFolder = 1;
 }
 
 def MatvecOp : LinalgStructured_Op<"matvec", [NInputs<2>, NOutputs<1>]> {
 

mlir/include/mlir/IR/AffineMap.h

public:
   using ImplType = detail::AffineMapStorage;
 
   constexpr AffineMap() : map(nullptr) {}
   explicit AffineMap(ImplType *map) : map(map) {}
 
   /// Returns a zero result affine map with no dimensions or symbols: () -> ().
   static AffineMap get(MLIRContext *context);
 
+  /// Returns a zero result affine map with `dimCount` dimensions and
+  /// `symbolCount` symbols, e.g.: `(...) -> ()`.
+  static AffineMap get(unsigned dimCount, unsigned symbolCount,
+                       MLIRContext *context);
+
   static AffineMap get(unsigned dimCount, unsigned symbolCount,
                        ArrayRef<AffineExpr> results);
 
   /// Returns a single constant result affine map.
   static AffineMap getConstantMap(int64_t val, MLIRContext *context);
 
   /// Returns an AffineMap with 'numDims' identity result dim exprs.
   static AffineMap getMultiDimIdentityMap(unsigned numDims,
inline raw_ostream &operator<<(raw_ostream &os, AffineMap map) {
   map.print(os);
   return os;
 }
 } // end namespace mlir
 
 namespace llvm {
 
 // AffineExpr hash just like pointers
-template <>
-struct DenseMapInfo<mlir::AffineMap> {
+template <> struct DenseMapInfo<mlir::AffineMap> {
   static mlir::AffineMap getEmptyKey() {
     auto pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
     return mlir::AffineMap(static_cast<mlir::AffineMap::ImplType *>(pointer));
   }
   static mlir::AffineMap getTombstoneKey() {
     auto pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
     return mlir::AffineMap(static_cast<mlir::AffineMap::ImplType *>(pointer));
   }

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

if (m.getNumSymbols() != 0)
       return op.emitOpError("expected indexing_map #")
              << idx << " to have no symbols";
 
     if (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

newSymbols[idx - numSymbolsThisMap] =
         getAffineSymbolExpr(idx, getContext());
   }
   auto newMap =
       map.replaceDimsAndSymbols(newDims, newSymbols, numDims, numSymbols);
   SmallVector<AffineExpr, 8> exprs;
   exprs.reserve(getResults().size());
   for (auto expr : getResults())
     exprs.push_back(expr.compose(newMap));
-  return AffineMap::get(numDims, numSymbols, exprs);
+  return exprs.empty() ? AffineMap::get(numDims, 0, map.getContext())
+                       : AffineMap::get(numDims, numSymbols, exprs);
 }
 
 bool AffineMap::isProjectedPermutation() {
   if (getNumSymbols() > 0)
     return false;
   SmallVector<bool, 8> seen(getNumInputs(), false);
   for (auto expr : getResults()) {
     if (auto dim = expr.dyn_cast<AffineDimExpr>()) {
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)
+  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) {
   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 results.empty() ? AffineMap::get(numDims, /*numSymbols=*/0,
+                                          maps.front().getContext())
+                         : AffineMap::get(numDims, /*numSymbols=*/0, results);
 }
 
 //===----------------------------------------------------------------------===//
 // MutableAffineMap.
 //===----------------------------------------------------------------------===//
 
 MutableAffineMap::MutableAffineMap(AffineMap map)
     : numDims(map.getNumDims()), numSymbols(map.getNumSymbols()),

mlir/lib/IR/MLIRContext.cpp

AffineMap AffineMap::getImpl(unsigned dimCount, unsigned symbolCount,
   });
 }
 
 AffineMap AffineMap::get(MLIRContext *context) {
   return getImpl(/*dimCount=*/0, /*symbolCount=*/0, /*results=*/{}, context);
 }
 
 AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
+                         MLIRContext *context) {
+  return getImpl(dimCount, /*symbolCount=*/0, /*results=*/{}, context);
+}
+
+AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
                          ArrayRef<AffineExpr> results) {
   // The number of results can't be zero.
   assert(!results.empty());
   return getImpl(dimCount, symbolCount, results, results[0].getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // Integer Sets: these are allocated into the bump pointer, and are immutable.

mlir/lib/Parser/Parser.cpp

AffineParser::parseAffineMapOfSSAIds(AffineMap &map,
   SmallVector<AffineExpr, 4> exprs;
   auto parseElt = [&]() -> ParseResult {
     auto elt = parseAffineExpr();
     exprs.push_back(elt);
     return elt ? success() : failure();
   };
 
   // Parse a multi-dimensional affine expression (a comma-separated list of
-  // 1-d affine expressions); the list cannot be empty. Grammar:
-  // multi-dim-affine-expr ::= `(` affine-expr (`,` affine-expr)* `)
+  // 1-d affine expressions); the list can be empty. Grammar:
+  // multi-dim-affine-expr ::= `(` `)`
+  //                         | `(` affine-expr (`,` affine-expr)* `)`
   if (parseCommaSeparatedListUntil(rightToken, parseElt,
                                    /*allowEmptyList=*/true))
     return failure();
   // Parsed a valid affine map.
   if (exprs.empty())
-    map = AffineMap::get(getContext());
+    map = AffineMap::get(numDimOperands, dimsAndSymbols.size() - numDimOperands,
+                         getContext());
   else
     map = AffineMap::get(numDimOperands, dimsAndSymbols.size() - numDimOperands,
                          exprs);
   return success();
 }
 
 /// Parse the range and sizes affine map definition inline.
 ///
AffineMap AffineParser::parseAffineMapRange(unsigned numDims,
   auto parseElt = [&]() -> ParseResult {
     auto elt = parseAffineExpr();
     ParseResult res = elt ? success() : failure();
     exprs.push_back(elt);
     return res;
   };
 
   // Parse a multi-dimensional affine expression (a comma-separated list of
-  // 1-d affine expressions); the list cannot be empty. Grammar:
-  // multi-dim-affine-expr ::= `(` affine-expr (`,` affine-expr)* `)
+  // 1-d affine expressions). Grammar:
+  // multi-dim-affine-expr ::= `(` `)`
+  //                         | `(` affine-expr (`,` affine-expr)* `)`
   if (parseCommaSeparatedListUntil(Token::r_paren, parseElt, true))
     return AffineMap();
 
   if (exprs.empty())
-    return AffineMap::get(getContext());
+    return AffineMap::get(numDims, numSymbols, getContext());
 
   // Parsed a valid affine map.
   return AffineMap::get(numDims, numSymbols, exprs);
 }
 
 /// Parse an affine constraint.
 ///  affine-constraint ::= affine-expr `>=` `0`
 ///                      | affine-expr `==` `0`

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<(i, j) -> ()>,
   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<(i) -> ()>
 ]
 
 #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<(i) -> ()>,
+  affine_map<(i) -> ()>
 ]
 
 #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<(i, j) -> ()>,
   affine_map<(i, j) -> (i, j)>
 ]
 
 #trait_broadcast = {
   args_in = 1,
   args_out = 1,
   indexing_maps = #broadcast_access,
   iterator_types = ["parallel", "parallel"],