[mlir][linalg] Allow TC ops taking an unused shaped operand.

[hanchung]

Summary

If one operand is not used in the formula, it will be considered a
shaped operand. And the result of indexing map of the operand will be the first
reduction dims.

Depends On D97383

[hanchung]

rebase

[hanchung]

fix and add more comments

[nicolasvasilache]

approving to unblock progress, this will be revisited soon in the context of the python eDSL

mlir/docs/Dialects/Linalg.md

1.  The input and output tensor parameters are specified as `id :
     general symbolic affine expressions.
 1.  The output shapes are specified explicitly, in TC they are always derived
     from the input shapes.
 1.  The operations used to specify computations use EDSC intrinsics so that they
     can easily be parsed and emitted into a simple region builder without
     resorting to more general MLIR parsing.
 1.  Reduction dimensions are specified with angle bracket notation on the
     operation they apply to (e.g. `std_add<k>` specifies that `k` is a reduction
-    dimension). In TC, a reduction is specified with `op=` operator and the
-    reduction dimensions are inferred.
+    dimension). In TC, the reduction dimensions are inferred. If one of the
+    operand is not used in any expressions, it will be considered a shape-only
+    operand, and the result of the indexing_map will be reduction dimensions.
 1.  The parallel and reduction dimension are ordered by the textual program
     order. For instance, in the comprehension `O(i, j) = std_add<k, l>(...)`,
     `i` (resp. `j`) is a parallel iterator encoded by affine dimension of
     position `0` (resp. `1`); `k` (resp. `l`) is a reduction iterator encoded by
     an affine dimension of position `2` (resp. `3`).
 1.  A list of attributes can be defined for the op with the format of `attr(
     strides: 2xi32)` and referenced in comprehension like `strides[0]`. These
     attribute uses will be parsed as affine symbols to generate op definition

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc

 //       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = std_subf([[d]], [[c]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 ods_def<Test8Op>:
 def test8(A: f32(M, K), B: f32(K)) -> (C: f32(M))
 {
   C(m) = std_subf<k>(std_mulf(A(m, k), B(k)), C(m));
 }
+
+// Test shape-only operand.
+// IMPL-LABEL:  ArrayAttr Test9Op::indexing_maps() {
+//       IMPL:    auto map0 = AffineMap::get(2, 2, {d0, d1}, context);
+//       IMPL:    auto map1 = AffineMap::get(2, 2, {d1}, context);
+//       IMPL:    auto map2 = AffineMap::get(2, 2, {d0}, context);
+ods_def<Test9Op>:
+def test9(A: f32(M, K), B: f32(K)) -> (C: f32(M))
+{
+  C(m) = std_addf<k>(C(m), A(m, k));
+}

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

visitPostorder(*pExpr, [&](const Expression &e) {
             "Unexpected multi-read of a tensor with different accesses");
         failed = true;
         return;
       }
       seenDefs.insert(use.tensorId);
       tensor.indexingMap = use.indexingMap;
       state.orderedTensorArgs[use] = tensor.index;
     });
-  state.numArgs = seenDefs.size();
+  // If more than one definitions are less. They are shaped-only operand, which
+  // are used to define reduction loops. For now, only accept exactly one
+  // shaped-only operand.
+  if (state.numArgs > seenDefs.size() + 1) {
+    failed = true;
+  } else if (state.numArgs == seenDefs.size() + 1) {
+    for (auto &tensorIter : registeredTensors) {
+      auto &tensor = tensorIter.getValue();
+      if (tensor.indexingMap)
+        continue;
+      if (auto *pTensorExpr =
+              dyn_cast<TensorExpr>(state.expressions[0].get())) {
+        SmallVector<AffineExpr, 4> exprs;
+        for (auto dim : pTensorExpr->reductionDimensions)
+          exprs.push_back(getAffineDimExpr(dim, parser.context));
+        tensor.indexingMap = AffineMap::get(state.dims.size(), symbols.size(),
+                                            exprs, parser.context);
+      }
+    }
+  }
   if (failed)
     return failure();
 
   return success();
 }
 
 /// Parse and print the information for a ODS def.
 ///
LogicalResult TCParser::parseAndEmitODSDef(llvm::raw_ostream &os) {
   }
 
   if (failed(parser.parseToken(Token::Kind::l_brace, "expected '{'")))
     return failure();
 
   SmallVector<ComprehensionParsingState, 4> perComprehensionStates;
   while (parser.curToken.isNot(Token::Kind::r_brace)) {
     perComprehensionStates.push_back(ComprehensionParsingState());
+    perComprehensionStates.back().numArgs = registeredTensors.size();
     if (failed(parseOneComprehension(cppOpName, tcName,
                                      perComprehensionStates.back())))
       return failure();
   };
   if (failed(parser.parseToken(Token::Kind::r_brace, "expected '}'")))
     return failure();
 
   // Print.
void TCParser::printReferenceIndexingMaps(llvm::raw_ostream &os,
   ss.flush();
 
   // 3. Get the list of affine maps for each input/output. The AffineExpr use
   // the common arithmetic operators on AffineExpr. These affine maps will
   // replace the `{2}` placeholder.
   std::string mapsStr;
   llvm::raw_string_ostream mapsStringStream(mapsStr);
 
-  SmallVector<TensorUse, 4> orderedUses(state.numArgs);
-  for (const auto &it : state.orderedTensorArgs)
-    orderedUses[it.second] = it.first;
-
   // Create a list of all symbols.
   SmallVector<std::string, 4> symbolReplacements;
   symbolReplacements.reserve(symbols.size());
   for (unsigned i = 0; i < symbols.size(); ++i) {
     const char *symFmt =
         "\n\tauto s{0} = getAffineSymbolExpr({0}, context); (void)s{0};";
     mapsStringStream << llvm::formatv(symFmt, i);
     symbolReplacements.push_back(llvm::formatv("s{0}", i));
const char *cstFmt =
         "\n\tauto cst{0} = getAffineConstantExpr({1}, context);";
     mapsStringStream << llvm::formatv(cstFmt, attrUse.index(), cstVal);
 
     unsigned position =
         attrUse.value().symbol.cast<AffineSymbolExpr>().getPosition();
     symbolReplacements[position] = llvm::formatv("cst{0}", attrUse.index());
   }
 
-  // For each tensor use, construct the affine map, replace symbols by the
-  // corresponding attribute values, and simplify the affine map.
-  for (auto tensorUse : llvm::enumerate(orderedUses)) {
-    auto indexingMap = tensorUse.value().indexingMap;
+  // For each registered tensor, construct the affine map, replace symbols by
+  // the corresponding attribute values, and simplify the affine map.
+  for (auto &tensorIter : registeredTensors) {
+    auto &tensor = tensorIter.getValue();
+    auto indexingMap = tensor.indexingMap;
     const char *mapFmt =
         "\n\tauto map{0} = AffineMap::get({1}, {2}, {3}, context);";
 
     std::string exprsStr;
     llvm::raw_string_ostream exprsStringStream(exprsStr);
     exprsStringStream << "{";
     llvm::interleaveComma(indexingMap.getResults(), exprsStringStream);
     exprsStringStream << "}";
     exprsStringStream.flush();
-    mapsStringStream << llvm::formatv(mapFmt, tensorUse.index(),
-                                      state.dims.size(),
+    mapsStringStream << llvm::formatv(mapFmt, tensor.index, state.dims.size(),
                                       indexingMap.getNumSymbols(), exprsStr);
 
     std::string replaceSymbolList =
         llvm::formatv("{ {0} }", llvm::join(symbolReplacements, ", "));
 
     // Note that we use `0` as the result affine map's number of symbols. All
     // symbols representing attribute usages should be folded away. But there
     // may exist additional symbols for tensor dimension upper bounds. Linalg
     // does not handle such cases right now. This needs to be fixed once we
     // need that.
     const char *replaceFmt =
         "\n\tmap{0} = map{0}.replaceDimsAndSymbols({{}, {1}, {2}, 0);";
-    mapsStringStream << llvm::formatv(replaceFmt, tensorUse.index(),
+    mapsStringStream << llvm::formatv(replaceFmt, tensor.index,
                                       replaceSymbolList, state.dims.size());
     const char *simplifyFmt = "\n\tmap{0} = simplifyAffineMap(map{0});";
-    mapsStringStream << llvm::formatv(simplifyFmt, tensorUse.index());
+    mapsStringStream << llvm::formatv(simplifyFmt, tensor.index);
   }
 
   mapsStringStream.flush();
 
   SmallVector<std::string, 4> mapList;
-  mapList.reserve(orderedUses.size());
-  for (unsigned i = 0; i < orderedUses.size(); ++i)
+  mapList.reserve(state.numArgs);
+  for (auto i : llvm::seq<unsigned>(0, state.numArgs))
     mapList.push_back(llvm::formatv("map{0}", i));
 
   // 4. Apply format to 1. using 2. and 3.
   os << llvm::formatv(referenceIndexingMapsFmt, cppOpName, dimsStr, mapsStr,
                       llvm::join(mapList, ", "));
 }
 
 /// Print the C++ StructuredOpsInterface impl of `regionBuilder`.