[Mlir-commits] [mlir] [mlir] Add helper to check elementwise-mappable ops with tensors and scalars (PR #154872)

[Adam Siemieniuk]

================
@@ -81,13 +105,39 @@ struct ConvertAnyElementwiseMappableOpOnRankedTensors : public RewritePattern {
       return rewriter.notifyMatchFailure(
           op, "requires elementwise op on ranked tensors");
 
-    auto rank = cast<RankedTensorType>(op->getResult(0).getType()).getRank();
-    SmallVector<AffineMap, 3> indexingMaps(
-        op->getNumResults() + op->getNumOperands(),
-        rewriter.getMultiDimIdentityMap(rank));
-    SmallVector<utils::IteratorType, 6> iteratorTypes(
+    auto resTy = cast<RankedTensorType>(op->getResult(0).getType());
+    auto rank = resTy.getRank();
+
+    // Maps: identity for tensors (rank > 0), scalar map for scalars/rank-0.
+    AffineMap scalarMap = AffineMap::get(/*dimCount=*/rank, /*symbolCount=*/0,
+                                         /*results=*/{}, rewriter.getContext());
+    AffineMap idMap = rewriter.getMultiDimIdentityMap(rank);
+
+    // Create indexing maps: one per operand, one per result.
+    SmallVector<AffineMap, 6> indexingMaps;
----------------
adam-smnk wrote:

nit: the sizes (like `, 6>` here and below) can be dropped

https://github.com/llvm/llvm-project/pull/154872