[Mlir-commits] [mlir] [mlir][Linalg] Allow more control in drop unit dims (PR #170104)

[Renato Golin]

================
@@ -535,6 +543,83 @@ struct ControlDropUnitDims {
     }
     return SmallVector<unsigned>{};
   };
+
+  using ComputeOperandShapeAndMapFnTy = std::function<UnitExtentReplacementInfo(
+      const ControlDropUnitDims &, MLIRContext *, IndexingMapOpInterface,
+      OpOperand *, DimensionMapping &, ArrayRef<AffineExpr>)>;
+  ComputeOperandShapeAndMapFnTy computeOperandShapeAndMapFn =
+      [](const ControlDropUnitDims &control, MLIRContext *context,
+         IndexingMapOpInterface op, OpOperand *opOperand,
+         DimensionMapping &oldDimsToNewDimsMap,
+         ArrayRef<AffineExpr> dimReplacements) -> UnitExtentReplacementInfo {
+    auto hasCollapsibleType = [](OpOperand &operand) {
+      Type operandType = operand.get().getType();
+      if (auto memrefOperandType = dyn_cast_or_null<MemRefType>(operandType)) {
+        return memrefOperandType.getLayout().isIdentity();
+      }
+      if (auto tensorOperandType = dyn_cast<RankedTensorType>(operandType)) {
+        return tensorOperandType.getEncoding() == nullptr;
+      }
+      return false;
+    };
+    auto indexingMap = op.getMatchingIndexingMap(opOperand);
+    SmallVector<int64_t> shape = op.getStaticOperandShape(opOperand);
+    if (!hasCollapsibleType(*opOperand)) {
----------------
rengolin wrote:

Nit: reverse the logic: `if (has...` for early exit.

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