[Mlir-commits] [mlir] [MLIR][Linalg] Introduce SpecializeOp (PR #70326)

[Aviad Cohen]

================
@@ -0,0 +1,52 @@
+//===- Specialize.cpp - linalg generic ops to named ops  ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a method to specialize generic operations to named
+// operations. Conceptually it is the opposite of generalize.cpp.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "linalg-specialization"
+
+using namespace mlir;
+using namespace mlir::linalg;
+
+static bool isaCopyOp(GenericOp genericOp) {
+  // Structural.
+  if (genericOp.getNumParallelLoops() != genericOp.getNumLoops())
+    return false;
+
+  // Operands and maps.
+  if (genericOp.getNumDpsInputs() != 1 || genericOp.getNumDpsInits() != 1)
+    return false;
+  auto mapRange = genericOp.getIndexingMapsArray();
+  if (mapRange.size() != 2 || !mapRange.front().isIdentity() ||
----------------
AviadCo wrote:

I came across AffineMaps which are actually identity maps if we take into consideration that the shape of the operands has 1 in it (this actually happens a lot when lowering from Tosa to Linalg), for example:
```
#map = affine_map<(d0, d1) -> (0, d1)>
Where operands shapes are <1x2xf32>
```

For those cases, I think it would be useful to have a util functions:

```

bool isIndexingMapsRepresentIdentity(linalg::GenericOp genericOp) {
  return llvm::all_of(
      genericOp->getOpOperands(), [&genericOp](OpOperand &opOperand) {
        return isCanonicalizedIdentityMap(
            genericOp.getMatchingIndexingMap(&opOperand), opOperand);
      });
}

} // namespace gcx::LinalgUtils

static bool isCanonicalizedIdentityMap(AffineMap map, OpOperand &opOperand) {
  if (map.getNumDims() != map.getNumResults())
    return false;
  ArrayRef<AffineExpr> results = map.getResults();
  auto shape = cast<ShapedType>(opOperand.get().getType()).getShape();
  for (unsigned i = 0, numDims = map.getNumDims(); i < numDims; ++i) {
    auto constExpr = results[i].dyn_cast<AffineConstantExpr>();
    if (constExpr && constExpr.getValue() == 0 && shape[0] == 1)
      continue;

    auto expr = results[i].dyn_cast<AffineDimExpr>();
    if (!expr || expr.getPosition() != i)
      return false;
  }
  return true;
}
```

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