[MLIR] Improve support for 0-dimensional Affine Maps.

Summary:
Modified AffineMap::get to remove support for the overload which allowed
an ArrayRef of AffineExpr but no context (and gathered the context from a
presumed first entry, resulting in bugs when there were 0 results).

Instead, we support only a ArrayRef and a context, and a version which
takes a single AffineExpr.

Additionally, removed some now needless case logic which previously
special cased which call to AffineMap::get to use.

Reviewers: flaub, bondhugula, rriddle!, nicolasvasilache, ftynse, ulysseB, mravishankar, antiagainst, aartbik

Subscribers: mehdi_amini, jpienaar, burmako, shauheen, antiagainst, arpith-jacob, mgester, lucyrfox, liufengdb, Joonsoo, bader, grosul1, frgossen, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D78226

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

@@ -188,9 +188,9 @@ def DotOp : LinalgStructured_Op<"dot", [NInputs<2>, NOutputs<1>]> {
       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}, context),
-        AffineMap::get(1, 0, {r_i}),
+        AffineMap::get(1, 0, {r_i}, context),
-        AffineMap::get(1, 0, context)};
+        AffineMap::get(1, 0, {}, context)};
     }
   }];
 
@@ -215,8 +215,10 @@ def MatvecOp : LinalgStructured_Op<"matvec", [NInputs<2>, NOutputs<1>]> {
       AffineExpr i, r_j;
       bindDims(context, i, r_j);
       return SmallVector<AffineMap, 8>{
-        AffineMap::get(2, 0, {i, r_j}), AffineMap::get(2, 0, {r_j}),
+        AffineMap::get(2, 0, {i, r_j}, context), 
-        AffineMap::get(2, 0, {i})};
+        AffineMap::get(2, 0, {r_j}, context),
+        AffineMap::get(2, 0, {i}, context)
+      };
     }
   }];
 
@@ -242,9 +244,11 @@ def MatmulOp : LinalgStructured_Op<"matmul", [NInputs<2>, NOutputs<1>]> {
       MLIRContext *context = getContext();
       AffineExpr i, j, r_k;
       bindDims(context, i, j, r_k);
-      return SmallVector<AffineMap, 8>{AffineMap::get(3, 0, {i, r_k}),
+      return SmallVector<AffineMap, 8>{
-                                       AffineMap::get(3, 0, {r_k, j}),
+        AffineMap::get(3, 0, {i, r_k}, context),
-                                       AffineMap::get(3, 0, {i, j})};
+        AffineMap::get(3, 0, {r_k, j},context),
+        AffineMap::get(3, 0, {i, j}, context)
+      };
     }
   }];
 
@@ -403,15 +407,15 @@ def ConvOp : PoolingBase_Op<"conv", [NInputs<2>, NOutputs<1>]> {
       auto ws = weightedPoolingInputIndex(*this, xs, zs);
       return SmallVector<AffineMap, 8>{
         // filter[z[0], ..., z[N-1], q, k]
-        AffineMap::get(idx, 0, concat(concat(zs, qs), ks)),
+        AffineMap::get(idx, 0, concat(concat(zs, qs), ks), context),
         // input[b,
         //       x[0]*s[0] + d[0]*z[0] - pad_low[0],
         //       ...
         //       x[N-1]*s[N-1] + d[N-1]*z[N-1] - pad_low[N-1],
         //       q]
-        AffineMap::get(idx, 0, concat(concat(bs, ws), qs)),
+        AffineMap::get(idx, 0, concat(concat(bs, ws), qs), context),
         // output[b, x[0], ..., x[N-1], k]
-        AffineMap::get(idx, 0, concat(concat(bs, xs), ks))};
+        AffineMap::get(idx, 0, concat(concat(bs, xs), ks), context)};
     }
   }];
 
@@ -465,11 +469,11 @@ class SingleInputPoolingBase_Op<string mnemonic>
           weightedPoolingInputIndex(*this, outputDims, windowDims);
       return SmallVector<AffineMap, 8>{
         // input
-        AffineMap::get(idx, 0, inputDims),
+        AffineMap::get(idx, 0, inputDims, context),
         // windowDims
-        AffineMap::get(idx, 0, windowDims),
+        AffineMap::get(idx, 0, windowDims, context),
         // output
-        AffineMap::get(idx, 0, outputDims)
+        AffineMap::get(idx, 0, outputDims, context)
         };
     }
   }];

mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h

@@ -25,22 +25,24 @@
 namespace mlir {
 
 inline bool isRowMajorMatmul(ArrayAttr indexingMaps) {
+  auto context = indexingMaps.getContext();
   AffineExpr m, n, k;
-  bindDims(indexingMaps.getContext(), m, n, k);
+  bindDims(context, m, n, k);
-  auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}));
+  auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}, context));
-  auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}));
+  auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}, context));
-  auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {m, n}));
+  auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {m, n}, context));
-  auto maps = ArrayAttr::get({mapA, mapB, mapC}, indexingMaps.getContext());
+  auto maps = ArrayAttr::get({mapA, mapB, mapC}, context);
   return indexingMaps == maps;
 }
 
 inline bool isColumnMajorMatmul(ArrayAttr indexingMaps) {
+  auto context = indexingMaps.getContext();
   AffineExpr m, n, k;
-  bindDims(indexingMaps.getContext(), m, n, k);
+  bindDims(context, m, n, k);
-  auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}));
+  auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}, context));
-  auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}));
+  auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}, context));
-  auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {n, m}));
+  auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {n, m}, context));
-  auto maps = ArrayAttr::get({mapA, mapB, mapC}, indexingMaps.getContext());
+  auto maps = ArrayAttr::get({mapA, mapB, mapC}, context);
   return indexingMaps == maps;
 }
 

mlir/include/mlir/IR/AffineMap.h

@@ -49,13 +49,13 @@ public:
   static AffineMap get(unsigned dimCount, unsigned symbolCount,
                        MLIRContext *context);
 
-  /// Returns an affine map with `dimCount` dimensions and `symbolCount` symbols
+  /// Returns an affine map with `dimCount` dimensions and `symbolCount` mapping
-  /// mapping to the given results. The array of results cannot be empty.
+  /// to a single output dimension
   static AffineMap get(unsigned dimCount, unsigned symbolCount,
-                       ArrayRef<AffineExpr> results);
+                       AffineExpr result);
 
   /// Returns an affine map with `dimCount` dimensions and `symbolCount` mapping
-  /// to the given results, where the number of results can be zero.
+  /// to the given results.
   static AffineMap get(unsigned dimCount, unsigned symbolCount,
                        ArrayRef<AffineExpr> results, MLIRContext *context);
 

mlir/lib/Analysis/AffineStructures.cpp

@@ -1464,11 +1464,8 @@ std::pair<AffineMap, AffineMap> FlatAffineConstraints::getLowerAndUpperBound(
     lbExprs.push_back(expr);
   }
 
-  auto lbMap = lbExprs.empty() ? AffineMap()
+  auto lbMap = AffineMap::get(dimCount, symCount, lbExprs, context);
-                               : AffineMap::get(dimCount, symCount, lbExprs);
+  auto ubMap = AffineMap::get(dimCount, symCount, ubExprs, context);
-
-  auto ubMap = ubExprs.empty() ? AffineMap()
-                               : AffineMap::get(dimCount, symCount, ubExprs);
 
   return {lbMap, ubMap};
 }

mlir/lib/Analysis/LoopAnalysis.cpp

@@ -62,8 +62,8 @@ void mlir::buildTripCountMapAndOperands(
 
   SmallVector<AffineExpr, 4> lbSplatExpr(ubValueMap.getNumResults(),
                                          lbMap.getResult(0));
-  auto lbMapSplat =
+  auto lbMapSplat = AffineMap::get(lbMap.getNumDims(), lbMap.getNumSymbols(),
-      AffineMap::get(lbMap.getNumDims(), lbMap.getNumSymbols(), lbSplatExpr);
+                                   lbSplatExpr, b.getContext());
   AffineValueMap lbSplatValueMap(lbMapSplat, forOp.getLowerBoundOperands());
 
   AffineValueMap tripCountValueMap;

mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp

@@ -99,11 +99,11 @@ SingleWorkgroupReduction::matchAsPerformingReduction(
       genericOp.indexing_maps().getValue()[1].cast<AffineMapAttr>();
   // The indexing map for the input should be `(i) -> (i)`.
   if (inputMap.getValue() !=
-      AffineMap::get(1, 0, {getAffineDimExpr(0, op->getContext())}))
+      AffineMap::get(1, 0, getAffineDimExpr(0, op->getContext())))
     return llvm::None;
   // The indexing map for the input should be `(i) -> (0)`.
   if (outputMap.getValue() !=
-      AffineMap::get(1, 0, {getAffineConstantExpr(0, op->getContext())}))
+      AffineMap::get(1, 0, getAffineConstantExpr(0, op->getContext())))
     return llvm::None;
 
   return linalg::RegionMatcher::matchAsScalarBinaryOp(genericOp);

mlir/lib/Dialect/Affine/EDSC/Builders.cpp

@@ -129,7 +129,7 @@ static ValueHandle createBinaryIndexHandle(
   if (v1) {
     operands.push_back(v1);
   }
-  auto map = AffineMap::get(numDims, numSymbols, {affCombiner(d0, d1)});
+  auto map = AffineMap::get(numDims, numSymbols, affCombiner(d0, d1));
   // TODO: createOrFold when available.
   Operation *op =
       makeComposedAffineApply(ScopedContext::getBuilder(),

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -522,7 +522,8 @@ AffineApplyNormalizer::AffineApplyNormalizer(AffineMap map,
          "Unexpected number of concatenated symbols");
   auto numDims = dimValueToPosition.size();
   auto numSymbols = concatenatedSymbols.size() - map.getNumSymbols();
-  auto auxiliaryMap = AffineMap::get(numDims, numSymbols, auxiliaryExprs);
+  auto auxiliaryMap =
+      AffineMap::get(numDims, numSymbols, auxiliaryExprs, map.getContext());
 
   LLVM_DEBUG(map.print(dbgs() << "\nCompose map: "));
   LLVM_DEBUG(auxiliaryMap.print(dbgs() << "\nWith map: "));
@@ -2163,19 +2164,13 @@ LogicalResult AffinePrefetchOp::fold(ArrayRef<Attribute> cstOperands,
 
 void AffineParallelOp::build(Builder *builder, OperationState &result,
                              ArrayRef<int64_t> ranges) {
-  // Default initialize empty maps.
-  auto lbMap = AffineMap::get(builder->getContext());
-  auto ubMap = AffineMap::get(builder->getContext());
-  // If there are ranges, set each to [0, N).
-  if (ranges.size()) {
   SmallVector<AffineExpr, 8> lbExprs(ranges.size(),
                                      builder->getAffineConstantExpr(0));
-    lbMap = AffineMap::get(0, 0, lbExprs);
+  auto lbMap = AffineMap::get(0, 0, lbExprs, builder->getContext());
   SmallVector<AffineExpr, 8> ubExprs;
   for (int64_t range : ranges)
     ubExprs.push_back(builder->getAffineConstantExpr(range));
-    ubMap = AffineMap::get(0, 0, ubExprs);
+  auto ubMap = AffineMap::get(0, 0, ubExprs, builder->getContext());
-  }
   build(builder, result, lbMap, {}, ubMap, {});
 }
 

mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp

@@ -51,8 +51,9 @@ void AffineValueMap::difference(const AffineValueMap &a,
     diffExprs.push_back(normalizer.getAffineMap().getResult(i) -
                         bMap.getResult(i));
 
-  auto diffMap = AffineMap::get(normalizer.getNumDims(),
+  auto diffMap =
-                                normalizer.getNumSymbols(), diffExprs);
+      AffineMap::get(normalizer.getNumDims(), normalizer.getNumSymbols(),
+                     diffExprs, aMap.getContext());
   canonicalizeMapAndOperands(&diffMap, &bOperands);
   diffMap = simplifyAffineMap(diffMap);
   res->reset(diffMap, bOperands);

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp

@@ -143,8 +143,9 @@ constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
       boundExprs.push_back(dim + tileSizes[i]);
       boundExprs.append(origUbMap.getResults().begin(),
                         origUbMap.getResults().end());
-      auto ubMap = AffineMap::get(origUbMap.getNumDims() + 1,
+      auto ubMap =
-                                  origUbMap.getNumSymbols(), boundExprs);
+          AffineMap::get(origUbMap.getNumDims() + 1, origUbMap.getNumSymbols(),
+                         boundExprs, b.getContext());
       newLoops[width + i].setUpperBound(/*operands=*/ubOperands, ubMap);
     } else {
       // No need of the min expression.

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

@@ -523,8 +523,10 @@ getSymbolLessAffineMaps(ArrayRef<ArrayRef<AffineExpr>> reassociation) {
          "Expected symbol-less expressions");
   SmallVector<AffineMap, 4> maps;
   maps.reserve(reassociation.size());
-  for (auto exprs : reassociation)
+  for (auto exprs : reassociation) {
-    maps.push_back(AffineMap::get(maxDim + 1, 0, exprs));
+    assert(exprs.size() != 0);
+    maps.push_back(AffineMap::get(maxDim + 1, 0, exprs, exprs[0].getContext()));
+  }
   return maps;
 }
 

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

@@ -252,7 +252,8 @@ public:
       // so having a max op is enough.
       auto maxMap = AffineMap::get(/*dimCount=*/1, 0,
                                    {getAffineDimExpr(/*position=*/0, context),
-                                    getAffineConstantExpr(0, context)});
+                                    getAffineConstantExpr(0, context)},
+                                   context);
       clampedImIdx.push_back(
           affine_max(dim.getType(), maxMap, ValueRange{dim}));
     }

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

@@ -294,7 +294,8 @@ makeTiledViews(OpBuilder &b, Location loc, LinalgOp linalgOp,
             /*dimCount=*/3, /*symbolCount=*/0,
             {getAffineDimExpr(/*position=*/0, b.getContext()),
              getAffineDimExpr(/*position=*/1, b.getContext()) -
-                 getAffineDimExpr(/*position=*/2, b.getContext())});
+                 getAffineDimExpr(/*position=*/2, b.getContext())},
+            b.getContext());
         auto d = folded_std_dim(folder, view, r);
         size = folded_affine_min(folder, b.getIndexType(), minMap,
                                  ValueRange{size, d, offset});

mlir/lib/Dialect/LoopOps/Transforms/ParallelLoopTiling.cpp

@@ -66,7 +66,8 @@ void mlir::loop::tileParallelLoop(ParallelOp op, ArrayRef<int64_t> tileSizes) {
       /*dimCount=*/3, /*symbolCount=*/0,
       {getAffineDimExpr(/*position=*/0, b.getContext()),
        getAffineDimExpr(/*position=*/1, b.getContext()) -
-           getAffineDimExpr(/*position=*/2, b.getContext())});
+           getAffineDimExpr(/*position=*/2, b.getContext())},
+      b.getContext());
 
   // Create the inner loop with adjusted bounds.
   SmallVector<Value, 2> newBounds;

mlir/lib/Dialect/Vector/VectorTransforms.cpp

@@ -1348,9 +1348,7 @@ private:
       auto targetExpr = getAffineDimExpr(idx < index ? idx : idx - 1, ctx);
       results.push_back(targetExpr);
     }
-    // The (...) -> () affine map has its own factory method.
+    return AffineMap::get(map.getNumDims() - 1, 0, results, ctx);
-    return results.empty() ? AffineMap::get(map.getNumDims() - 1, 0, ctx)
-                           : AffineMap::get(map.getNumDims() - 1, 0, results);
   }
 
   // Helper to drop dimension from vector type.

mlir/lib/Dialect/Vector/VectorUtils.cpp

@@ -176,7 +176,7 @@ static AffineMap makePermutationMap(
            "Vectorization prerequisite violated: at most 1 index may be "
            "invariant wrt a vectorized loop");
   }
-  return AffineMap::get(indices.size(), 0, perm);
+  return AffineMap::get(indices.size(), 0, perm, context);
 }
 
 /// Implementation detail that walks up the parents and records the ones with

mlir/lib/IR/AffineMap.cpp

@@ -104,7 +104,7 @@ AffineMap AffineMap::getPermutationMap(ArrayRef<unsigned> permutation,
   for (auto index : permutation)
     affExprs.push_back(getAffineDimExpr(index, context));
   auto m = std::max_element(permutation.begin(), permutation.end());
-  auto permutationMap = AffineMap::get(*m + 1, 0, affExprs);
+  auto permutationMap = AffineMap::get(*m + 1, 0, affExprs, context);
   assert(permutationMap.isPermutation() && "Invalid permutation vector");
   return permutationMap;
 }
@@ -127,13 +127,16 @@ static void getMaxDimAndSymbol(ArrayRef<AffineExprContainer> exprsList,
 template <typename AffineExprContainer>
 static SmallVector<AffineMap, 4>
 inferFromExprList(ArrayRef<AffineExprContainer> exprsList) {
+  assert(!exprsList.empty());
+  assert(!exprsList[0].empty());
+  auto context = exprsList[0][0].getContext();
   int64_t maxDim = -1, maxSym = -1;
   getMaxDimAndSymbol(exprsList, maxDim, maxSym);
   SmallVector<AffineMap, 4> maps;
   maps.reserve(exprsList.size());
   for (const auto &exprs : exprsList)
     maps.push_back(AffineMap::get(/*dimCount=*/maxDim + 1,
-                                  /*symbolCount=*/maxSym + 1, exprs));
+                                  /*symbolCount=*/maxSym + 1, exprs, context));
   return maps;
 }
 
@@ -153,7 +156,7 @@ AffineMap AffineMap::getMultiDimIdentityMap(unsigned numDims,
   dimExprs.reserve(numDims);
   for (unsigned i = 0; i < numDims; ++i)
     dimExprs.push_back(mlir::getAffineDimExpr(i, context));
-  return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs);
+  return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs, context);
 }
 
 MLIRContext *AffineMap::getContext() const { return map->context; }
@@ -255,8 +258,7 @@ AffineMap AffineMap::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
     results.push_back(
         expr.replaceDimsAndSymbols(dimReplacements, symReplacements));
 
-  return results.empty() ? get(numResultDims, 0, getContext())
+  return get(numResultDims, numResultSyms, results, getContext());
-                         : get(numResultDims, numResultSyms, results);
 }
 
 AffineMap AffineMap::compose(AffineMap map) {
@@ -280,8 +282,7 @@ AffineMap AffineMap::compose(AffineMap map) {
   exprs.reserve(getResults().size());
   for (auto expr : getResults())
     exprs.push_back(expr.compose(newMap));
-  return exprs.empty() ? AffineMap::get(numDims, 0, map.getContext())
+  return AffineMap::get(numDims, numSymbols, exprs, map.getContext());
-                       : AffineMap::get(numDims, numSymbols, exprs);
 }
 
 bool AffineMap::isProjectedPermutation() {
@@ -312,7 +313,7 @@ AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) {
   for (auto idx : resultPos) {
     exprs.push_back(getResult(idx));
   }
-  return AffineMap::get(getNumDims(), getNumSymbols(), exprs);
+  return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());
 }
 
 AffineMap mlir::simplifyAffineMap(AffineMap map) {
@@ -321,7 +322,8 @@ AffineMap mlir::simplifyAffineMap(AffineMap map) {
     exprs.push_back(
         simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols()));
   }
-  return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs);
+  return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs,
+                        map.getContext());
 }
 
 AffineMap mlir::removeDuplicateExprs(AffineMap map) {
@@ -354,7 +356,7 @@ AffineMap mlir::inversePermutation(AffineMap map) {
       seenExprs.push_back(expr);
   if (seenExprs.size() != map.getNumInputs())
     return AffineMap();
-  return AffineMap::get(map.getNumResults(), 0, seenExprs);
+  return AffineMap::get(map.getNumResults(), 0, seenExprs, map.getContext());
 }
 
 AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
@@ -369,9 +371,8 @@ AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
     results.append(m.getResults().begin(), m.getResults().end());
     numDims = std::max(m.getNumDims(), numDims);
   }
-  return results.empty() ? AffineMap::get(numDims, /*numSymbols=*/0,
+  return AffineMap::get(numDims, /*numSymbols=*/0, results,
-                                          maps.front().getContext())
+                        maps.front().getContext());
-                         : AffineMap::get(numDims, /*numSymbols=*/0, results);
 }
 
 //===----------------------------------------------------------------------===//
@@ -380,8 +381,7 @@ AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
 
 MutableAffineMap::MutableAffineMap(AffineMap map)
     : numDims(map.getNumDims()), numSymbols(map.getNumSymbols()),
-      // A map always has at least 1 result by construction
+      context(map.getContext()) {
-      context(map.getResult(0).getContext()) {
   for (auto result : map.getResults())
     results.push_back(result);
 }
@@ -390,8 +390,7 @@ void MutableAffineMap::reset(AffineMap map) {
   results.clear();
   numDims = map.getNumDims();
   numSymbols = map.getNumSymbols();
-  // A map always has at least 1 result by construction
+  context = map.getContext();
-  context = map.getResult(0).getContext();
   for (auto result : map.getResults())
     results.push_back(result);
 }
@@ -416,5 +415,5 @@ void MutableAffineMap::simplify() {
 }
 
 AffineMap MutableAffineMap::getAffineMap() const {
-  return AffineMap::get(numDims, numSymbols, results);
+  return AffineMap::get(numDims, numSymbols, results, context);
 }

mlir/lib/IR/Builders.cpp

@@ -293,12 +293,11 @@ AffineMap Builder::getEmptyAffineMap() { return AffineMap::get(context); }
 
 AffineMap Builder::getConstantAffineMap(int64_t val) {
   return AffineMap::get(/*dimCount=*/0, /*symbolCount=*/0,
-                        {getAffineConstantExpr(val)});
+                        getAffineConstantExpr(val));
 }
 
 AffineMap Builder::getDimIdentityMap() {
-  return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0,
+  return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, getAffineDimExpr(0));
-                        {getAffineDimExpr(0)});
 }
 
 AffineMap Builder::getMultiDimIdentityMap(unsigned rank) {
@@ -306,18 +305,19 @@ AffineMap Builder::getMultiDimIdentityMap(unsigned rank) {
   dimExprs.reserve(rank);
   for (unsigned i = 0; i < rank; ++i)
     dimExprs.push_back(getAffineDimExpr(i));
-  return AffineMap::get(/*dimCount=*/rank, /*symbolCount=*/0, dimExprs);
+  return AffineMap::get(/*dimCount=*/rank, /*symbolCount=*/0, dimExprs,
+                        context);
 }
 
 AffineMap Builder::getSymbolIdentityMap() {
   return AffineMap::get(/*dimCount=*/0, /*symbolCount=*/1,
-                        {getAffineSymbolExpr(0)});
+                        getAffineSymbolExpr(0));
 }
 
 AffineMap Builder::getSingleDimShiftAffineMap(int64_t shift) {
   // expr = d0 + shift.
   auto expr = getAffineDimExpr(0) + shift;
-  return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, {expr});
+  return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr);
 }
 
 AffineMap Builder::getShiftedAffineMap(AffineMap map, int64_t shift) {
@@ -325,7 +325,8 @@ AffineMap Builder::getShiftedAffineMap(AffineMap map, int64_t shift) {
   shiftedResults.reserve(map.getNumResults());
   for (auto resultExpr : map.getResults())
     shiftedResults.push_back(resultExpr + shift);
-  return AffineMap::get(map.getNumDims(), map.getNumSymbols(), shiftedResults);
+  return AffineMap::get(map.getNumDims(), map.getNumSymbols(), shiftedResults,
+                        context);
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/IR/MLIRContext.cpp

@@ -717,10 +717,8 @@ AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
 }
 
 AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
-                         ArrayRef<AffineExpr> results) {
+                         AffineExpr result) {
-  // The number of results can't be zero.
+  return getImpl(dimCount, symbolCount, {result}, result.getContext());
-  assert(!results.empty());
-  return getImpl(dimCount, symbolCount, results, results[0].getContext());
 }
 
 AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,

mlir/lib/IR/StandardTypes.cpp

@@ -723,7 +723,7 @@ MemRefType mlir::canonicalizeStridedLayout(MemRefType t) {
       simplifyAffineExpr(m.getResult(0), m.getNumDims(), m.getNumSymbols());
   if (expr != simplifiedLayoutExpr)
     return MemRefType::Builder(t).setAffineMaps({AffineMap::get(
-        m.getNumDims(), m.getNumSymbols(), {simplifiedLayoutExpr})});
+        m.getNumDims(), m.getNumSymbols(), simplifiedLayoutExpr)});
   return MemRefType::Builder(t).setAffineMaps({});
 }
 

mlir/lib/Parser/Parser.cpp

@@ -3132,12 +3132,8 @@ AffineParser::parseAffineMapOfSSAIds(AffineMap &map,
                                    /*allowEmptyList=*/true))
     return failure();
   // Parsed a valid affine map.
-  if (exprs.empty())
   map = AffineMap::get(numDimOperands, dimsAndSymbols.size() - numDimOperands,
-                         getContext());
+                       exprs, getContext());
-  else
-    map = AffineMap::get(numDimOperands, dimsAndSymbols.size() - numDimOperands,
-                         exprs);
   return success();
 }
 
@@ -3166,11 +3162,8 @@ AffineMap AffineParser::parseAffineMapRange(unsigned numDims,
   if (parseCommaSeparatedListUntil(Token::r_paren, parseElt, true))
     return AffineMap();
 
-  if (exprs.empty())
-    return AffineMap::get(numDims, numSymbols, getContext());
-
   // Parsed a valid affine map.
-  return AffineMap::get(numDims, numSymbols, exprs);
+  return AffineMap::get(numDims, numSymbols, exprs, getContext());
 }
 
 /// Parse an affine constraint.

mlir/lib/Transforms/LoopFusion.cpp

@@ -942,7 +942,8 @@ static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
   }
   auto indexRemap = zeroOffsetCount == rank
                         ? AffineMap()
-                        : AffineMap::get(outerIVs.size() + rank, 0, remapExprs);
+                        : AffineMap::get(outerIVs.size() + rank, 0, remapExprs,
+                                         forOp.getContext());
   // Replace all users of 'oldMemRef' with 'newMemRef'.
   LogicalResult res =
       replaceAllMemRefUsesWith(oldMemRef, newMemRef, {}, indexRemap,

mlir/lib/Transforms/PipelineDataTransfer.cpp

@@ -98,8 +98,8 @@ static bool doubleBuffer(Value oldMemRef, AffineForOp forOp) {
   // Create 'iv mod 2' value to index the leading dimension.
   auto d0 = bInner.getAffineDimExpr(0);
   int64_t step = forOp.getStep();
-  auto modTwoMap = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0,
+  auto modTwoMap =
-                                  {d0.floorDiv(step) % 2});
+      AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, d0.floorDiv(step) % 2);
   auto ivModTwoOp = bInner.create<AffineApplyOp>(forOp.getLoc(), modTwoMap,
                                                  forOp.getInductionVar());
 

mlir/lib/Transforms/Utils/LoopUtils.cpp

@@ -103,7 +103,8 @@ static void getCleanupLoopLowerBound(AffineForOp forOp, unsigned unrollFactor,
   operands.clear();
   operands.push_back(lb);
   operands.append(bumpValues.begin(), bumpValues.end());
-  map = AffineMap::get(1 + tripCountMap.getNumResults(), 0, newUbExprs);
+  map = AffineMap::get(1 + tripCountMap.getNumResults(), 0, newUbExprs,
+                       b.getContext());
   // Simplify the map + operands.
   fullyComposeAffineMapAndOperands(&map, &operands);
   map = simplifyAffineMap(map);
@@ -485,7 +486,7 @@ LogicalResult mlir::loopUnrollByFactor(AffineForOp forOp,
     if (!forOpIV.use_empty()) {
       // iv' = iv + 1/2/3...unrollFactor-1;
       auto d0 = builder.getAffineDimExpr(0);
-      auto bumpMap = AffineMap::get(1, 0, {d0 + i * step});
+      auto bumpMap = AffineMap::get(1, 0, d0 + i * step);
       auto ivUnroll =
           builder.create<AffineApplyOp>(forOp.getLoc(), bumpMap, forOpIV);
       operandMap.map(forOpIV, ivUnroll);
@@ -616,7 +617,7 @@ LogicalResult mlir::loopUnrollJamByFactor(AffineForOp forOp,
       if (!forOpIV.use_empty()) {
         // iv' = iv + i, i = 1 to unrollJamFactor-1.
         auto d0 = builder.getAffineDimExpr(0);
-        auto bumpMap = AffineMap::get(1, 0, {d0 + i * step});
+        auto bumpMap = AffineMap::get(1, 0, d0 + i * step);
         auto ivUnroll =
             builder.create<AffineApplyOp>(forOp.getLoc(), bumpMap, forOpIV);
         operandMap.map(forOpIV, ivUnroll);
@@ -859,7 +860,8 @@ static void augmentMapAndBounds(OpBuilder &b, Value iv, AffineMap *map,
   auto bounds = llvm::to_vector<4>(map->getResults());
   bounds.push_back(b.getAffineDimExpr(map->getNumDims()) + offset);
   operands->insert(operands->begin() + map->getNumDims(), iv);
-  *map = AffineMap::get(map->getNumDims() + 1, map->getNumSymbols(), bounds);
+  *map = AffineMap::get(map->getNumDims() + 1, map->getNumSymbols(), bounds,
+                        b.getContext());
   canonicalizeMapAndOperands(map, operands);
 }
 
@@ -1514,7 +1516,7 @@ generatePointWiseCopy(Location loc, Value memref, Value fastMemRef,
     b = forOp.getBodyBuilder();
 
     auto fastBufOffsetMap =
-        AffineMap::get(lbOperands.size(), 0, {fastBufOffsets[d]});
+        AffineMap::get(lbOperands.size(), 0, fastBufOffsets[d]);
     auto offset = b.create<AffineApplyOp>(loc, fastBufOffsetMap, lbOperands);
 
     // Construct the subscript for the fast memref being copied into/from:
@@ -1529,7 +1531,8 @@ generatePointWiseCopy(Location loc, Value memref, Value fastMemRef,
     memIndices.push_back(forOp.getInductionVar());
   }
 
-  auto fastBufMap = AffineMap::get(2 * rank, /*symbolCount=*/0, fastBufExprs);
+  auto fastBufMap =
+      AffineMap::get(2 * rank, /*symbolCount=*/0, fastBufExprs, b.getContext());
   fullyComposeAffineMapAndOperands(&fastBufMap, &fastBufMapOperands);
   fastBufMap = simplifyAffineMap(fastBufMap);
   canonicalizeMapAndOperands(&fastBufMap, &fastBufMapOperands);
@@ -1837,7 +1840,8 @@ static LogicalResult generateCopy(
     auto dimExpr = b.getAffineDimExpr(regionSymbols.size() + i);
     remapExprs.push_back(dimExpr - fastBufOffsets[i]);
   }
-  auto indexRemap = AffineMap::get(regionSymbols.size() + rank, 0, remapExprs);
+  auto indexRemap = AffineMap::get(regionSymbols.size() + rank, 0, remapExprs,
+                                   b.getContext());
 
   // Record the begin since it may be invalidated by memref replacement.
   Block::iterator prevOfBegin;

mlir/test/Dialect/Affine/simplify-affine-structures.mlir

@@ -271,3 +271,13 @@ func @affine.apply(%N : index) {
   // CHECK-NEXT: addi
   return
 }
+
+// -----
+
+// CHECK-DAG: #[[MAP_0D:.*]] = affine_map<() -> ()>
+
+// CHECK-LABEL: func @simplify_zero_dim_map
+func @simplify_zero_dim_map(%in : memref<f32>) -> f32 {
+  %out = affine.load %in[] : memref<f32>
+  return %out : f32
+}