[mlir][Linalg] Add pass to remove unit-extent dims from tensor operands of Generic ops.

[mravishankar]

Summary

Unit-extent dimensions are typically used for achieving broadcasting
behavior. The pattern added (along with canonicalization patterns
added previously) removes the use of unit-extent dimensions, and
instead uses a more canonical representation of the computation. This
new pattern is not added as a canonicalization for now since it
entails adding additional reshape operations. A pass is added to
exercise these patterns, along with an API entry to populate a
patterns list with these patterns.

[nicolasvasilache]

Generally looks good, thanks Mahesh.
I left some code laering comments.

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

74

This does not belong here right?
It also intersects with this revision: https://reviews.llvm.org/D79759
Skipping since I am not seeing it used in this revision.

82

In the grander scheme, should this function be exposed?

I am asking because the way it is currently used, it takes and indexMap that comes from a GenericOp.
Are there cases where we'd want to use this but the genericOp may be invalid IR (because inversePermutation would fail) ?

Should this be refactored so that it can also help build correct IR that is canonical by construction?

No need to necessarily change, I am mostly asking form your higher-level user's perspective.

94

auto zero = getAffineConstantExpr(0, context);

then just: return shape[fdim] == 1 && exprs[dim] == zero;.

Why can dim overflow here?
IMO, it would be more natural to while (dim < rank && isUnitExtent(dim))

99

This should assert at least one dim is not a unitExtent otherwise reassociationMaps is empty?

101

trivial braces here and below.

104

I find the multi-dim++ a bit tricky to follow but I can't see offhand how to make it more readable.

172

typo: result

1142

The layering does not seem right to me. Traditionally we either:

1. create an op canonicalization patterns for a specific op (i.e. set hasCanonicalization = 1 in Ops.td) and later collect them with OpName::getCanonicalizationPatterns() but this is not what you want here.
2. what you want here, create a new pass or transform and put everything there this way Ops.cpp does not ned to include Passes.h

If you're worried that other things are not visible they should also be moved in the proper place.
Your populate impl should resemble:

patterns.insert<ReplaceUnitExtentTensors>(context);
GenericOp::getCanonicalizationPatterns(patterns);
TensorReshapeOp::getCanonicalizationPatterns(patterns);

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

918

This should be its own pass in its own file, with its own pattern it is unclear to me why it is in the Fusion pass.

[mravishankar]

Folding with previous change in stack and addressing comments.

[mravishankar]

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

74

Ah, sorry. This shouldnt be here . Removed.

82

I am not sure what you mean by exposed. This is a static function that is useful only for the callsite. The issue here might be that this was added to LinalgOps.cpp. Moving all of this into a separate pass. So it is not visible to op parsing/verification, etc. So in effect it isnt exposed outside of the pass.

Regarding other questions

1. It should work on any indexMap. I dont see this to be related to whether the genericOp is valid or not since it is only looking at a single index map. The validity of the genericOp using the indexMap constructed here is the concern of the caller AFAICS
2. I dont know how to refactor to help build correct IR. I cant trivially see how each index map construction of a generic op can be steered for the final generic op to be valid. The validity depends on concatenation and inversion. How do I break it down to the level of individual index maps?

94

Changed. Thanks!

99

Turns out it is OK to have all dims being unitExtent. Added a test for that.

104

Made this a while loop and made the dim++ more explicit.

1142

I agree. I moved everything to a separate file and added the relevant patterns into that.

[nicolasvasilache]

Thanks!

Note the last minor points.

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

101

It seems like this whole function could be replaced by:

1. create an affine map that either sets unitDims to 0 or drops them (depending on whether you want a projection or explictly set to 0).
2. compose with the indexingMaps

?

143

We have ArrayAttr::getAsRange, seems a similar ArrayAttr::getFromRange would be useful?
No need to do it in this CL if this is too intrusive.

176

I am curious why does the iterator type matter here?

193

nit:

for (auto attr : llvm::enumerate(iteratorTypes))
  if (unitDims.count(attr.index()) == 0)
    newIteratorTypes.push_back(attr.value());

?

216

typo: dimension

244

Thanks, much simpler to follow AFAIR !

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

166

nit:newline

[rriddle]

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

114

nit: Drop trivial braces.

250

nit: ++dim;

256

Same here.

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

578–579

nit: Merge this into the if and use a type other than auto.

[mravishankar]

ddressing comments

[mravishankar]

Fixing a bug in the code.

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

InterfaceMethod<
       "Return the number of inputs, irrespective of their buffer or tensor "
       "type, and output buffers",
       "unsigned", "getNumInputsAndOutputBuffers"
     >,
     InterfaceMethod<
       "Return the range over inputs (irrespective of type) and output buffers.",
       "Operation::operand_range", "getInputsAndOutputBuffers"
     >,
+    InterfaceMethod<
+      "Return the shaped types for all the inputs and outputs",
+      "SmallVector<ShapedType, 4>", "getInputOutputShapedTypes"
+    >,
 
     //===------------------------------------------------------------------===//
     // Other interface methods.
     //===------------------------------------------------------------------===//
     InterfaceMethod<
       "Return the reference iterators for this named op (if any are "
       "specified). These reference iterators are used to specify the default "
       "behavior of the op. Typically this would be a static method but in "
let methods = [
     InterfaceMethod<
       "Return the iterator types attribute within the current operation.",
       "ArrayAttr", "iterator_types"
     >,
     InterfaceMethod<
       "Return the indexing maps attribute within the current operation.",
       "ArrayAttr", "indexing_maps"
     >,
+    InterfaceMethod<
+      "Return the indexing maps within the current operation.",
+      "SmallVector<AffineMap, 4>", "getIndexingMaps"
+    >,
     InterfaceMethod<"Return the input or output indexing map at index `i`.",
       "AffineMap", "getIndexingMap", (ins "unsigned":$i)
     >,
     InterfaceMethod<"Return the input indexing map at index `i`.",
       "AffineMap", "getInputIndexingMap", (ins "unsigned":$i)
     >,
     InterfaceMethod<"Return the output indexing map at index `i`.",
       "AffineMap", "getOutputIndexingMap", (ins "unsigned":$i)

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

public:
   ShapedType getShapedType(unsigned i) {
     if (i < nInputs())
       return getInputShapedType(i);
     if (i < getNumInputsAndOutputBuffers())
       return getOutputBufferType(i - nInputs()).template cast<ShapedType>();
     return getOutputTensorTypes()[i - getNumInputsAndOutputBuffers()]
         .template cast<ShapedType>();
   }
+  /// Return the shaped types for all the inputs and outputs
+  SmallVector<ShapedType, 4> getInputOutputShapedTypes() {
+    SmallVector<Type, 4> inputOutputTypes(
+        this->getOperation()->operand_type_begin(),
+        this->getOperation()->operand_type_end());
+    inputOutputTypes.append(this->getOperation()->result_type_begin(),
+                            this->getOperation()->result_type_end());
+    return llvm::to_vector<4>(
+        llvm::map_range(inputOutputTypes, [](Type type) -> ShapedType {
+          return type.cast<ShapedType>();
+        }));
+  }
 
   //==========================================================================//
   // Other interface methods.
   //==========================================================================//
 
   // Get or build the indexing_maps ArrayAttr.
   ArrayAttr iterator_types() {
     // Return the attribute if it is present.
ArrayAttr indexing_maps() {
     SmallVector<Attribute, 4> attrs{attrRange.begin(), attrRange.end()};
     auto attr = ArrayAttr::get(attrs, ctx);
     // TODO(ntv): Need to memoize this. Can't just store as an attribute atm as
     // it will impact parser, printer and tests.
     // this->getOperation()->setAttr("indexing_maps", attr);
     return attr;
   }
 
+  SmallVector<AffineMap, 4> getIndexingMaps() {
+    return llvm::to_vector<4>(
+        llvm::map_range(indexing_maps(), [](Attribute attr) -> AffineMap {
+          return attr.cast<AffineMapAttr>().getValue();
+        }));
+  }
+
   AffineMap getIndexingMap(unsigned i) {
     assert(i < getNumInputsAndOutputs());
     return indexing_maps()
         .getValue()[i]
         .template cast<AffineMapAttr>()
         .getValue();
   }
 

mlir/include/mlir/Dialect/Linalg/Passes.h

 namespace mlir {
 class FuncOp;
 class MLIRContext;
 class ModuleOp;
 template <typename T> class OperationPass;
 class OwningRewritePatternList;
 class Pass;
 
+std::unique_ptr<OperationPass<FuncOp>> createLinalgFoldUnitExtentDimsPass();
+
 std::unique_ptr<OperationPass<FuncOp>> createLinalgFusionPass();
 std::unique_ptr<Pass> createLinalgFusionOfTensorOpsPass();
 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingPass(ArrayRef<int64_t> tileSizes = {});
 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingToParallelLoopsPass(ArrayRef<int64_t> tileSizes = {});

 /// buffers instead.
 std::unique_ptr<OperationPass<ModuleOp>>
 createConvertLinalgOnTensorsToBuffersPass();
 
 /// Patterns for fusing linalg operation on tensors.
 void populateLinalgTensorOpsFusionPatterns(MLIRContext *context,
                                            OwningRewritePatternList &patterns);
 
+/// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
+/// tensors.
+void populateLinalgFoldUnitExtentDimsPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_PASSES_H_

mlir/include/mlir/Dialect/Linalg/Passes.td

 //===-- Passes.td - Linalg pass definition file ------------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LINALG_PASSES
 #define MLIR_DIALECT_LINALG_PASSES
 
 include "mlir/Pass/PassBase.td"
 
+def LinalgFoldUnitExtentDims : FunctionPass<"linalg-fold-unit-extent-dims"> {
+  let summary = "Remove unit-extent dimension in Linalg ops on tensors";
+  let constructor = "mlir::createLinalgFoldUnitExtentDimsPass()";
+  let options = [
+    Option<"foldOneTripLoopsOnly", "fold-one-trip-loops-only", "bool",
+            /*default=*/"false",
+      "Only folds the one-trip loops from Linalg ops on tensors "
+      "(for testing purposes only)">
+  ];
+}
+
 def LinalgFusion : FunctionPass<"linalg-fusion"> {
   let summary = "Fuse operations in the linalg dialect";
   let constructor = "mlir::createLinalgFusionPass()";
 }
 
 def LinalgFusionOfTensorOps : Pass<"linalg-fusion-for-tensor-ops"> {
   let summary = "Fuse operations on RankedTensorType in linalg dialect";
   let constructor = "mlir::createLinalgFusionOfTensorOpsPass()";

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

 // LinalgOps.td), we define an overloaded `print` function and a
 // parse`className` function.
 
 //===----------------------------------------------------------------------===//
 // GenericOps
 //===----------------------------------------------------------------------===//
 
 template <typename GenericOpType>
 static void printGenericOp(OpAsmPrinter &p, GenericOpType op) {
   auto attrNames = op.linalgTraitAttrNames();
   llvm::StringSet<> linalgTraitAttrsSet;
   linalgTraitAttrsSet.insert(attrNames.begin(), attrNames.end());
   SmallVector<NamedAttribute, 8> attrs;
   for (auto attr : op.getAttrs())
     if (linalgTraitAttrsSet.count(attr.first.strref()) > 0)
       attrs.push_back(attr);
 
   auto dictAttr = DictionaryAttr::get(attrs, op.getContext());
   p << op.getOperationName() << " " << dictAttr;
   p.printOptionalAttrDict(op.getAttrs(), attrNames);
   p << " " << op.getOperands();
   if (!op.region().empty())
     p.printRegion(op.region());
   p << ": " << op.getOperandTypes();
   auto outputTensorTypes = op.getResultTypes();
   if (!outputTensorTypes.empty())
     p << " -> " << outputTensorTypes;
 }
 
 static void print(OpAsmPrinter &p, GenericOp op) { printGenericOp(p, op); }
 
 static void print(OpAsmPrinter &p, IndexedGenericOp op) {
   printGenericOp(p, op);
 }
 
 static ParseResult parseGenericOp(OpAsmParser &parser, OperationState &result) {
   SmallVector<OpAsmParser::OperandType, 8> operandsInfo, regionOperandsInfo;
   DictionaryAttr dictAttr;
   // Parse the core linalg traits that must check into a dictAttr.
   // The name is unimportant as we will overwrite result.attributes.
   // The core linalg traits must contain the information necessary to pass the
   // verifier.
   if (parser.parseAttribute(dictAttr, "_", result.attributes))
     return failure();
   result.attributes.assign(dictAttr.getValue().begin(),
                            dictAttr.getValue().end());
 
LogicalResult BlockArgsVerifier<IndexedGenericOp>::verify(IndexedGenericOp op,
                                                           Block &block) {
   auto nInputViews = op.getNumInputs();
   auto nLoops = op.getNumLoops();
   auto nOperands = op.getNumOperands();
   if (block.getNumArguments() != nOperands + nLoops)
     return op.emitOpError(
         "expected number of block arguments to match number of operands + "
         "number of loops");
 
   // Note: the number and type of yield values are checked in the YieldOp.
   for (unsigned i = 0; i < nLoops; ++i)
     if (!block.getArgument(i).getType().isIndex())
       return op.emitOpError("expected block argument ")
              << (i + 1) << " to be an index";
 
   for (unsigned i = 0; i < nOperands; ++i) {
     unsigned memrefArgIndex = i + nLoops;

 ///
 /// is folded into
 ///
 ///   result = [affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>,
 ///             affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>]
 static ArrayAttr collapseReassociationMaps(ArrayRef<AffineMap> mapsProducer,
                                            ArrayRef<AffineMap> mapsConsumer,
                                            MLIRContext *context) {
-  if (mapsProducer.size() == 0 || mapsConsumer.size() == 0 ||
+  if (mapsProducer.empty() || mapsConsumer.empty() ||
       mapsProducer[0].getNumDims() < mapsConsumer[0].getNumDims() ||
       mapsProducer.size() != mapsConsumer[0].getNumDims())
     return nullptr;
   unsigned numLhsDims = mapsProducer[0].getNumDims();
   unsigned currDim = 0;
   SmallVector<AffineExpr, 4> reassociations;
   SmallVector<Attribute, 4> reassociationMaps;
   for (AffineMap rhs : mapsConsumer) {
     for (AffineExpr rhsExpr : rhs.getResults()) {
       AffineDimExpr dimExpr = rhsExpr.cast<AffineDimExpr>();
       for (int i = 0, e = mapsProducer[dimExpr.getPosition()].getNumResults();
-           i != e; ++i) {
+           i < e; ++i) {
         reassociations.push_back(getAffineDimExpr(currDim++, context));
       }
     }
     reassociationMaps.push_back(AffineMapAttr::get(AffineMap::get(
         numLhsDims, /*numSymbols =*/0, reassociations, context)));
     reassociations.clear();
   }
   return ArrayAttr::get(reassociationMaps, context);
OpFoldResult ReshapeOp::fold(ArrayRef<Attribute>) {
   return foldReshapeOp(*this);
 }
 OpFoldResult SliceOp::fold(ArrayRef<Attribute>) {
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
   return {};
 }
 OpFoldResult TensorReshapeOp::fold(ArrayRef<Attribute>) {
-  if (succeeded(foldMemRefCast(*this)))
-    return getResult();
   return foldReshapeOp(*this);
 }
 OpFoldResult TransposeOp::fold(ArrayRef<Attribute>) {
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // Auto-generated Linalg named ops.
 //===----------------------------------------------------------------------===//
 
 template <typename NamedStructuredOpType>
 void buildNamedStructuredOpRegionAndAttributes(Builder &builder,
                                                OperationState &result,
                                                TypeRange operandTypes,
                                                TypeRange tensorResultTypes) {
   Region &region = *result.addRegion();
   Block *body = new Block();

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRLinalgTransforms
+  DropUnitDims.cpp
   Fusion.cpp
   Interchange.cpp
   Loops.cpp
   Promotion.cpp
   TensorsToBuffers.cpp
   Tiling.cpp
   Transforms.cpp
   Vectorization.cpp

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

+//===- DropUnitDims.cpp - Pass to drop use of unit-extent for broadcasting ===//
+//
+// 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 patterns/pass to remove usage of unit-extent dimensions
+// to specify broadcasting in favor of more canonical representation of the
+// computation
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
+#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/FoldUtils.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "linalg-drop-unit-dims"
+
+using namespace mlir;
+using namespace mlir::edsc;
+using namespace mlir::edsc::intrinsics;
+using namespace mlir::linalg;
+
+/// Implements a pass that canonicalizes the uses of unit-extent dimensions for
+/// broadcasting. For example,
+///
+/// ```mlir
+/// #accesses = [
+///   affine_map<(d0, d1) -> (0, d1)>,
+///   affine_map<(d0, d1) -> (d0, 0)>,
+///   affine_map<(d0, d1) -> (d0, d1)>
+/// ]
+///
+/// #trait = {
+///   args_in = 2,
+///   args_out = 1,
+///   indexing_maps = #accesses,
+///   iterator_types = ["parallel", "parallel"],
+///   library_call = "some_external_fn"
+/// }
+///
+/// func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) ->
+/// tensor<5x5xf32>
+/// {
+///   %0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1) -> (d0, d1)>] :
+///        tensor<5xf32> into tensor<1x5xf32>
+///   %1 = linalg.tensor_reshape %arg1 [affine_map<(d0, d1) -> (d0, d1)>] :
+///        tensor<5xf32> into tensor<5x1xf32>
+///   %2 = linalg.generic #trait %0, %1 {
+///        ^bb0(%arg2: f32, %arg3: f32):
+///          %3 = addf %arg2, %arg3 : f32
+///          linalg.yield %3 : f32
+///        } : tensor<1x5xf32>, tensor<5x1xf32> -> tensor<5x5xf32>
+///   return %2 : tensor<5x5xf32>
+/// }
+///
+/// would canonicalize to
+///
+/// ```mlir
+/// #accesses = [
+///   affine_map<(d0, d1) -> (d1)>,
+///   affine_map<(d0, d1) -> (d0)>,
+///   affine_map<(d0, d1) -> (d0, d1)>
+/// ]
+///
+/// #trait = {
+///   args_in = 2,
+///   args_out = 1,
+///   indexing_maps = #accesses,
+///   iterator_types = ["parallel", "parallel"],
+///   library_call = "some_external_fn"
+/// }
+///
+/// func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) ->
+/// tensor<5x5xf32>
+/// {
+///   %0 = linalg.generic #trait %arg0, %arg1 {
+///        ^bb0(%arg2: f32, %arg3: f32):
+///          %3 = addf %arg2, %arg3 : f32
+///          linalg.yield %3 : f32
+///        } : tensor<5xf32>, tensor<5xf32> -> tensor<5x5xf32>
+///   return %0 : tensor<5x5xf32>
+/// }
+
+/// Given dims of the iteration space of a structured op that are known to be
+/// single trip count (`unitDims`), return the indexing maps to use in the
+/// canonicalized op with these dims removed, given the original `indexingMaps`.
+static ArrayAttr replaceUnitDims(DenseSet<unsigned> &unitDims,
+                                 ArrayRef<AffineMap> indexingMaps,
+                                 MLIRContext *context) {
+  if (indexingMaps.empty())
+    return nullptr;
+  unsigned numIterationDims = indexingMaps.front().getNumDims();
+  unsigned numSymbols = indexingMaps.front().getNumSymbols();
+
+  // Compute the replacement for each dim expr.
+  SmallVector<AffineExpr, 4> dimReplacements;
+  dimReplacements.reserve(numIterationDims);
+  unsigned numKeptDims = 0;
+  for (unsigned dim : llvm::seq<unsigned>(0, numIterationDims)) {
+    if (unitDims.count(dim))
+      dimReplacements.push_back(getAffineConstantExpr(0, context));
+    else
+      dimReplacements.push_back(getAffineDimExpr(numKeptDims++, context));
+  }
+
+  // Symbols remain the same.
+  SmallVector<AffineExpr, 4> symReplacements;
+  symReplacements.reserve(numSymbols);
+  for (unsigned symbol : llvm::seq<unsigned>(0, numSymbols))
+    symReplacements.push_back(getAffineSymbolExpr(symbol, context));
+
+  SmallVector<AffineMap, 4> newIndexingMaps;
+  newIndexingMaps.reserve(indexingMaps.size());
+  for (AffineMap operandMap : indexingMaps) {
+    // Expected indexing maps to have no symbols.
+    if (operandMap.getNumSymbols())
+      return nullptr;
+    newIndexingMaps.push_back(simplifyAffineMap(
+        operandMap.replaceDimsAndSymbols(dimReplacements, symReplacements,
+                                         numIterationDims - unitDims.size(),
+                                         numSymbols)));
+  }
+
+  // Check that the new index maps are invertible. If not, something went
+  // wrong, so abort.
+  if (!inversePermutation(concatAffineMaps(newIndexingMaps)))
+    return nullptr;
+  return ArrayAttr::get(
+      llvm::to_vector<4>(llvm::map_range(
+          newIndexingMaps,
+          [](AffineMap map) -> Attribute { return AffineMapAttr::get(map); })),
+      context);
+}
+
+namespace {
+/// Pattern to fold unit-trip count loops in GenericOps.
+// TODO: Generalize this to indexed-generic as well by modifying the region args
+// as well.
+struct FoldUnitDimLoops : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(GenericOp genericOp,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<AffineMap, 4> indexingMaps = genericOp.getIndexingMaps();
+    if (indexingMaps.empty())
+      return failure();
+
+    // Check if any of the iteration dimensions are unit-trip count. They will
+    // end up being unit-trip count if they are used to index into a unit-dim
+    // tensor/memref.
+    AffineMap invertedMap = inversePermutation(concatAffineMaps(indexingMaps));
+    if (!invertedMap)
+      return failure();
+    SmallVector<int64_t, 4> dims;
+    for (ShapedType shapedType : genericOp.getInputOutputShapedTypes())
+      dims.append(shapedType.getShape().begin(), shapedType.getShape().end());
+    DenseSet<unsigned> unitDims;
+    ArrayAttr iteratorTypes = genericOp.iterator_types();
+    for (auto expr : enumerate(invertedMap.getResults())) {
+      if (AffineDimExpr dimExpr = expr.value().dyn_cast<AffineDimExpr>())
+        if (dims[dimExpr.getPosition()] == 1 &&
+            iteratorTypes[expr.index()].dyn_cast<StringAttr>().getValue() ==
+                getParallelIteratorTypeName())
+          unitDims.insert(expr.index());
+    }
+    if (unitDims.empty())
+      return failure();
+
+    // Compute the modified indexing maps.
+    MLIRContext *context = rewriter.getContext();
+    ArrayAttr newIndexingMapAttr =
+        replaceUnitDims(unitDims, indexingMaps, context);
+    if (!newIndexingMapAttr)
+      return genericOp.emitError("unable to compute modified indexing_maps");
+
+    // Compute the iterator types of the modified op by dropping the one-trip
+    // count loops.
+    SmallVector<Attribute, 4> newIteratorTypes;
+    for (auto attr : llvm::enumerate(iteratorTypes)) {
+      if (!unitDims.count(attr.index()))
+        newIteratorTypes.push_back(attr.value());
+    }
+
+    rewriter.startRootUpdate(genericOp);
+    genericOp.indexing_mapsAttr(newIndexingMapAttr);
+    genericOp.iterator_typesAttr(ArrayAttr::get(newIteratorTypes, context));
+    rewriter.finalizeRootUpdate(genericOp);
+    return success();
+  }
+};
+
+struct UnitExtentReplacementInfo {
+  RankedTensorType type;
+  AffineMap indexMap;
+  ArrayAttr reassociation;
+};
+} // namespace
+
+/// Utility function for replacing operands/results to a linalg generic
+/// operation on tensors with unit-extent dimensions. These can be replaced with
+/// an operand/result with the unit-extent dimension removed. This is only done
+/// if the indexing map used to access that didimensionmension has a
+/// AffineConstantExpr of value 0. Given the `type` of an result/operand of a
+/// Linalg op, and its `indexMap` the utility function returns:
+/// - the new type with dimensions of size 1 removed.
+/// - modified index map that can be used to access the replaced result/operand
+/// - the reassociation that converts from the original tensor type to the
+///   modified tensor type.
+static UnitExtentReplacementInfo replaceUnitExtents(AffineMap indexMap,
+                                                    RankedTensorType type,
+                                                    MLIRContext *context) {
+  ArrayRef<int64_t> shape = type.getShape();
+  ArrayRef<AffineExpr> exprs = indexMap.getResults();
+  SmallVector<AffineExpr, 2> reassociations;
+  SmallVector<Attribute, 4> reassociationMaps;
+  SmallVector<AffineExpr, 4> newIndexExprs;
+  SmallVector<int64_t, 4> newShape;
+
+  int64_t origRank = type.getRank();
+  AffineExpr zeroExpr = getAffineConstantExpr(0, context);
+  auto isUnitExtent = [&](int64_t dim) -> bool {
+    return shape[dim] == 1 && exprs[dim] == zeroExpr;
+  };
+
+  unsigned dim = 0;
+  // Fold dimensions that are unit-extent at the beginning of the tensor.
+  while (dim < origRank && isUnitExtent(dim))
+    reassociations.push_back(getAffineDimExpr(dim++, context));
+  while (dim < origRank) {
+    reassociations.push_back(getAffineDimExpr(dim, context));
+    newIndexExprs.push_back(exprs[dim]);
+    newShape.push_back(shape[dim]);
+    // Fold all following dimensions that are unit-extent.
+    while (dim + 1 < origRank && isUnitExtent(dim + 1)) {
+      ++dim;
+      reassociations.push_back(getAffineDimExpr(dim, context));
+    }
+    reassociationMaps.push_back(AffineMapAttr::get(AffineMap::get(
+        origRank, /*numSymbols = */ 0, reassociations, context)));
+    reassociations.clear();
+    ++dim;
+  }
+  UnitExtentReplacementInfo info = {
+      RankedTensorType::get(newShape, type.getElementType()),
+      AffineMap::get(indexMap.getNumDims(), indexMap.getNumSymbols(),
+                     newIndexExprs, context),
+      ArrayAttr::get(reassociationMaps, context)};
+  return info;
+}
+
+namespace {
+/// Pattern to replace tensors operands/results that are unit extents.
+struct ReplaceUnitExtentTensors : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(GenericOp genericOp,
+                                PatternRewriter &rewriter) const override {
+    if (!genericOp.hasTensorSemantics())
+      return failure();
+
+    MLIRContext *context = rewriter.getContext();
+    Location loc = genericOp.getLoc();
+
+    SmallVector<AffineMap, 4> newIndexingMaps;
+    SmallVector<ArrayAttr, 4> reassociationMaps;
+    SmallVector<ShapedType, 4> newInputOutputTypes;
+    bool doCanonicalization = false;
+    for (auto it : llvm::zip(genericOp.getIndexingMaps(),
+                             genericOp.getInputOutputShapedTypes())) {
+      auto replacementInfo = replaceUnitExtents(
+          std::get<0>(it), std::get<1>(it).cast<RankedTensorType>(), context);
+      reassociationMaps.push_back(replacementInfo.reassociation);
+      newIndexingMaps.push_back(replacementInfo.indexMap);
+      newInputOutputTypes.push_back(replacementInfo.type);
+      doCanonicalization =
+          doCanonicalization || replacementInfo.type != std::get<1>(it);
+    }
+
+    // If the indexing maps of the result operation are not invertible (i.e. not
+    // legal), abort.
+    if (!doCanonicalization ||
+        !inversePermutation(concatAffineMaps(newIndexingMaps)))
+      return failure();
+
+    // If any operand type change, insert a reshape to convert from the original
+    // type to the new type.
+    SmallVector<Value, 4> newOperands;
+    newOperands.reserve(genericOp.getNumOperands());
+    for (auto operand : llvm::enumerate(genericOp.getOperands())) {
+      if (operand.value().getType() == newInputOutputTypes[operand.index()]) {
+        newOperands.push_back(operand.value());
+      } else {
+        newOperands.push_back(rewriter.create<linalg::TensorReshapeOp>(
+            loc, newInputOutputTypes[operand.index()], operand.value(),
+            reassociationMaps[operand.index()]));
+      }
+    }
+
+    // If any result type change, insert a reshape to convert from the original
+    // type to the new type.
+    SmallVector<Type, 4> resultTypes;
+    resultTypes.reserve(genericOp.getNumResults());
+    for (unsigned i : llvm::seq<unsigned>(0, genericOp.getNumResults()))
+      resultTypes.push_back(
+          newInputOutputTypes[i + genericOp.getNumOperands()]);
+    GenericOp replacementOp = rewriter.create<GenericOp>(
+        loc, resultTypes, newOperands, genericOp.args_in(),
+        genericOp.args_out(), rewriter.getAffineMapArrayAttr(newIndexingMaps),
+        genericOp.iterator_types(),
+        /*doc = */ nullptr,
+        /*library_call = */ nullptr);
+    rewriter.inlineRegionBefore(genericOp.region(), replacementOp.region(),
+                                replacementOp.region().begin());
+
+    // If any result tensor has a modified shape, then add reshape to recover
+    // the original shape.
+    SmallVector<Value, 4> resultReplacements;
+    for (auto result : llvm::enumerate(replacementOp.getResults())) {
+      unsigned index = result.index() + replacementOp.getNumOperands();
+      RankedTensorType origResultType = genericOp.getResult(result.index())
+                                            .getType()
+                                            .cast<RankedTensorType>();
+      if (origResultType != result.value().getType()) {
+        resultReplacements.push_back(rewriter.create<linalg::TensorReshapeOp>(
+            loc, origResultType, result.value(), reassociationMaps[index]));
+      } else {
+        resultReplacements.push_back(result.value());
+      }
+    }
+    rewriter.replaceOp(genericOp, resultReplacements);
+    return success();
+  }
+};
+} // namespace
+
+/// Patterns that are used to canonicalize the use of unit-extent dims for
+/// broadcasting.
+void mlir::populateLinalgFoldUnitExtentDimsPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns) {
+  patterns.insert<FoldUnitDimLoops, ReplaceUnitExtentTensors>(context);
+  TensorReshapeOp::getCanonicalizationPatterns(patterns, context);
+}
+
+namespace {
+/// Pass that removes unit-extent dims within generic ops.
+struct LinalgFoldUnitExtentDimsPass
+    : public LinalgFoldUnitExtentDimsBase<LinalgFoldUnitExtentDimsPass> {
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    FuncOp funcOp = getFunction();
+    MLIRContext *context = funcOp.getContext();
+    if (foldOneTripLoopsOnly)
+      patterns.insert<FoldUnitDimLoops>(context);
+    else
+      populateLinalgFoldUnitExtentDimsPatterns(context, patterns);
+    applyPatternsAndFoldGreedily(funcOp.getBody(), patterns);
+  }
+};
+} // namespace
+
+std::unique_ptr<OperationPass<FuncOp>>
+mlir::createLinalgFoldUnitExtentDimsPass() {
+  return std::make_unique<LinalgFoldUnitExtentDimsPass>();
+}

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

static void generateFusedRegion(PatternRewriter &rewriter,
     }
 
     // Add operations from producer (except the yield operation) to the fused
     // op.
     for (auto &op : producerBlock.getOperations()) {
       if (auto yieldOp = dyn_cast<YieldOp>(op)) {
         // Lookup the value the yield operation is mapped to.
         Value yieldVal = yieldOp.getOperand(0);
-        auto clonedVal = mapper.lookup(yieldVal);
+        if (Value clonedVal = mapper.lookupOrNull(yieldVal))
           mapper.map(consumerBlock.getArgument(consumerIdx), clonedVal);
         continue;
       }
       rewriter.clone(op, mapper);
     }
     for (auto &op : consumerBlock.getOperations())
       rewriter.clone(op, mapper);
   }
 };

 } // namespace
 
 void mlir::populateLinalgTensorOpsFusionPatterns(
     MLIRContext *context, OwningRewritePatternList &patterns) {
   patterns.insert<FuseTensorOps<GenericOp>, FuseTensorOps<TensorReshapeOp>>(
       context);
 }
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgFusionPass() {
   return std::make_unique<LinalgFusionPass>();
 }
 
 std::unique_ptr<Pass> mlir::createLinalgFusionOfTensorOpsPass() {
   return std::make_unique<FusionOfTensorOpsPass>();
 }

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

+// RUN: mlir-opt %s -linalg-fold-unit-extent-dims -split-input-file | FileCheck %s
+
+#accesses = [
+  affine_map<(i, j, k, l, m) -> (i, k, m)>,
+  affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>
+]
+
+#trait = {
+  args_in = 1,
+  args_out = 1,
+  iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],
+  indexing_maps = #accesses,
+  library_call = "some_external_func"
+}
+
+func @drop_one_trip_loops(%arg0 : tensor<?x1x?xf32>) -> tensor<?x1x?x1x?xf32>
+{
+  %0 = linalg.generic #trait %arg0 {
+       ^bb0(%arg1 : f32) :
+         linalg.yield %arg1 : f32
+       } : tensor<?x1x?xf32> -> tensor<?x1x?x1x?xf32>
+  return %0 : tensor<?x1x?x1x?xf32>
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+//   CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1, d2) -> (d2)>
+//   CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+//   CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+//   CHECK-DAG: #[[MAP4:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>
+//   CHECK-DAG: #[[MAP5:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3)>
+//   CHECK-DAG: #[[MAP6:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d4)>
+// CHECK-LABEL: func @drop_one_trip_loops
+//       CHECK: linalg.tensor_reshape %{{.*}} [#[[MAP0]], #[[MAP1]]]
+//       CHECK: linalg.generic
+//  CHECK-SAME:   indexing_maps = [#[[MAP2]], #[[MAP3]]]
+//  CHECK-SAME:   iterator_types = ["parallel", "parallel", "parallel"]
+//       CHECK: linalg.tensor_reshape %{{.*}} [#[[MAP4]], #[[MAP5]], #[[MAP6]]]
+
+// -----
+
+#map0 = affine_map<(i, j) -> (i, j)>
+#access = [#map0, #map0]
+#trait = {
+  args_in = 1,
+  args_out = 1,
+  iterator_types = ["parallel", "parallel"],
+  indexing_maps = #access,
+  library_call = "some_external_func"
+}
+
+func @drop_all_loops(%arg0 : tensor<1x1xf32>) -> tensor<1x1xf32>
+{
+  %0 = linalg.generic #trait %arg0 {
+       ^bb0(%arg1: f32) :
+         linalg.yield %arg1 : f32
+       } : tensor<1x1xf32> -> tensor<1x1xf32>
+  return %0 : tensor<1x1xf32>
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<() -> ()>
+// CHECK-LABEL: func @drop_all_loops
+//       CHECK:   linalg.tensor_reshape %{{.*}} []
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP0]]]
+//  CHECK-SAME:     iterator_types = []
+
+// -----
+
+#accesses = [
+  affine_map<(d0) -> (0, d0)>,
+  affine_map<(d0) -> (d0)>
+]
+
+#trait = {
+  args_in = 1,
+  args_out = 1,
+  indexing_maps = #accesses,
+  iterator_types = ["parallel"],
+  library_call = "some_external_fn"
+}
+
+func @leading_dim_1_canonicalization(%arg0: tensor<1x5xf32>) -> tensor<5xf32> {
+  %0 = linalg.generic #trait %arg0 {
+  ^bb0(%arg2: f32):     // no predecessors
+    linalg.yield %arg2 : f32
+  }  : tensor<1x5xf32> -> tensor<5xf32>
+  return %0 : tensor<5xf32>
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func @leading_dim_1_canonicalization
+//       CHECK:   linalg.tensor_reshape %{{.*}} [#[[MAP0]]]
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     indexing_maps = [#[[MAP1]], #[[MAP1]]]
+//  CHECK-SAME:     iterator_types = ["parallel"]
+
+// -----
+
+#accesses = [
+  affine_map<(d0, d1) -> (0, d1)>,
+  affine_map<(d0, d1) -> (d0, 0)>,
+  affine_map<(d0, d1) -> (d0, d1)>
+]
+
+#trait = {
+  args_in = 2,
+  args_out = 1,
+  indexing_maps = #accesses,
+  iterator_types = ["parallel", "parallel"],
+  library_call = "some_external_fn"
+}
+
+func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) -> tensor<5x5xf32>
+{
+  %0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1) -> (d0, d1)>] :
+       tensor<5xf32> into tensor<1x5xf32>
+  %1 = linalg.tensor_reshape %arg1 [affine_map<(d0, d1) -> (d0, d1)>] :
+       tensor<5xf32> into tensor<5x1xf32>
+  %2 = linalg.generic #trait %0, %1 {
+       ^bb0(%arg2: f32, %arg3: f32):
+         %3 = addf %arg2, %arg3 : f32
+         linalg.yield %3 : f32
+       } : tensor<1x5xf32>, tensor<5x1xf32> -> tensor<5x5xf32>
+  return %2 : tensor<5x5xf32>
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1) -> (d1)>
+//   CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d0)>
+//   CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func @broadcast_test
+//   CHECK-NOT:   linalg.tensor_reshape
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+//  CHECK-SAME:     iterator_types = ["parallel", "parallel"]
+//   CHECK-NOT:   linalg.tensor_reshape
+
+// -----
+
+#accesses = [
+  affine_map<(d0, d1) -> (0, 0)>,
+  affine_map<(d0, d1) -> (d0, d1)>
+]
+
+#trait = {
+  args_in = 1,
+  args_out = 1,
+  indexing_maps = #accesses,
+  iterator_types = ["parallel", "parallel"],
+  library_call = "some_external_fn"
+}
+
+func @broadcast_scalar(%arg0 : tensor<1x1xf32>) -> tensor<?x?xf32>
+{
+   %0 = linalg.generic #trait %arg0 {
+        ^bb0(%arg1 : f32):
+    linalg.yield %arg1 : f32
+  } : tensor<1x1xf32> -> tensor<?x?xf32>
+   return %0 : tensor<?x?xf32>
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1) -> ()>
+//   CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func @broadcast_scalar
+//  CHECK-SAME:   %[[ARG0:.*]]: tensor<1x1xf32>
+//       CHECK:   %[[A:.*]] = linalg.tensor_reshape %[[ARG0]] []
+//  CHECK-SAME:     tensor<1x1xf32> into tensor<f32>
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
+//  CHECK-SAME:     iterator_types = ["parallel", "parallel"]
+//  CHECK-SAME:     %[[A]]

mlir/test/Dialect/Linalg/fold-unit-trip-loops.mlir

+// RUN: mlir-opt %s -linalg-fold-unit-extent-dims="fold-one-trip-loops-only" -split-input-file | FileCheck %s
+
+#accesses = [
+  affine_map<(i, j, k, l, m) -> (i, k, m)>,
+  affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>
+]
+
+#trait = {
+  args_in = 1,
+  args_out = 1,
+  iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],
+  indexing_maps = #accesses,
+  library_call = "some_external_func"
+}
+
+func @drop_one_trip_loops(%arg0 : tensor<?x1x?xf32>) -> tensor<?x1x?x1x?xf32>
+{
+  %0 = linalg.generic #trait %arg0 {
+       ^bb0(%arg1 : f32) :
+         linalg.yield %arg1 : f32
+       } : tensor<?x1x?xf32> -> tensor<?x1x?x1x?xf32>
+  return %0 : tensor<?x1x?x1x?xf32>
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d2)>
+//   CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d1, 0, d2)>
+// CHECK-LABEL: func @drop_one_trip_loops
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
+//  CHECK-SAME:     iterator_types = ["parallel", "parallel", "parallel"]
+
+// -----
+
+#map0 = affine_map<(i, j) -> (i, j)>
+#access = [#map0, #map0]
+#trait = {
+  args_in = 1,
+  args_out = 1,
+  iterator_types = ["parallel", "parallel"],
+  indexing_maps = #access,
+  library_call = "some_external_func"
+}
+
+func @drop_all_loops(%arg0 : tensor<1x1xf32>) -> tensor<1x1xf32>
+{
+  %0 = linalg.generic #trait %arg0 {
+       ^bb0(%arg1: f32) :
+         linalg.yield %arg1 : f32
+       } : tensor<1x1xf32> -> tensor<1x1xf32>
+  return %0 : tensor<1x1xf32>
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<() -> (0, 0)>
+// CHECK-LABEL: func @drop_all_loops
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP0]]]
+//  CHECK-SAME:     iterator_types = []
+
+// -----
+
+#map0 = affine_map<(i, j) -> (i, j)>
+#access = [#map0, #map0]
+#trait = {
+  args_in = 1,
+  args_out = 1,
+  iterator_types = ["parallel", "parallel"],
+  indexing_maps = #access,
+  library_call = "some_external_func"
+}
+
+func @drop_all_loops(%arg0 : memref<1x1xf32>, %arg1 : memref<1x1xf32>)
+{
+  linalg.generic #trait %arg0, %arg1 {
+    ^bb0(%arg2: f32, %arg3 : f32) :
+      linalg.yield %arg2 : f32
+    } : memref<1x1xf32>, memref<1x1xf32>
+  return
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<() -> (0, 0)>
+// CHECK-LABEL: func @drop_all_loops
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP0]]]
+//  CHECK-SAME:     iterator_types = []
+
+// -----
+
+#accesses = [
+  affine_map<(d0, d1) -> (d0, d1)>,
+  affine_map<(d0, d1) -> (d1)>
+]
+
+#trait = {
+  args_in = 1,
+  args_out = 1,
+  indexing_maps = #accesses,
+  iterator_types = ["parallel", "parallel"],
+  library_call = "some_external_fn"
+}
+
+func @leading_dim_1_canonicalization(%arg0: tensor<1x5xf32>) -> tensor<5xf32> {
+  %0 = linalg.generic #trait %arg0 {
+  ^bb0(%arg2: f32):     // no predecessors
+    linalg.yield %arg2 : f32
+  }  : tensor<1x5xf32> -> tensor<5xf32>
+  return %0 : tensor<5xf32>
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0) -> (0, d0)>
+//   CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: func @leading_dim_1_canonicalization
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
+//  CHECK-SAME:     iterator_types = ["parallel"]