[mlir] Move `linalg.tensor_expand/collapse_shape` to TensorDialect.

[pifon2a]

Summary

RFC: https://llvm.discourse.group/t/rfc-reshape-ops-restructuring/3310

linalg.fill gets a canonicalizer, because FoldFillWithTensorReshape cannot be moved to tensorops (it uses linalg::FillOp inside). Before it was listed as a canonicalization pattern for the reshape operations, now it became a canonicalization for FillOp.

[pifon2a]

Clean-up

[nicolasvasilache]

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

221

IIUC you need this to change the anchor point of canonicalization?
Could you please specify it in the commit message?

[pifon2a]

Thank you! Added a comment.

[nicolasvasilache]

Awesome, thanks much @pifon2a !

mlir/include/mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h

31

Thanks for spelling this out, this has been an issue for some time and will need to be addressed but in the meantime this is a good step forward.
Thanks for putting in the effort to untangle cycles despite the pain!

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

551–554

I understand this reanchoring can be changed once we have a tensor dialect op that replaces FillOp ?
@gysit could you please make a note of this?

[mehdi_amini]

mlir/include/mlir/InitAllDialects.h

88

I'm very concerned by this. This looks like a big hack, did you check with someone like River before doing something like this in the core infra?

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

 /// other AffineApplyOps supplying those operands. The operands of the resulting
 /// AffineApplyOp do not change the length of  AffineApplyOp chains.
 AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map,
                                       ValueRange operands);
 /// Variant of `makeComposedAffineApply` which infers the AffineMap from `e`.
 AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineExpr e,
                                       ValueRange values);
 
+/// Returns the values obtained by applying `map` to the list of values.
+SmallVector<Value, 4> applyMapToValues(OpBuilder &b, Location loc,
+                                       AffineMap map, ValueRange values);
+
 /// Given an affine map `map` and its input `operands`, this method composes
 /// into `map`, maps of AffineApplyOps whose results are the values in
 /// `operands`, iteratively until no more of `operands` are the result of an
 /// AffineApplyOp. When this function returns, `map` becomes the composed affine
 /// map, and each Value in `operands` is guaranteed to be either a loop IV or a
 /// terminal symbol, i.e., a symbol defined at the top level or a block/function
 /// argument.
 void fullyComposeAffineMapAndOperands(AffineMap *map,

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

 namespace linalg {
 class LinalgOp;
 
 /// OpOperand vector that implicitly converts to a Value vector.
 struct OpOperandVector : public SmallVector<OpOperand *> {
   operator SmallVector<Value>();
 };
 
-/// Returns the values obtained by applying `map` to the list of values.
-SmallVector<Value, 4> applyMapToValues(OpBuilder &b, Location loc,
-                                       AffineMap map, ValueRange values);
-
 /// Checks whether `linalgOp` conforms to ContractionOpInterface.
 // TODO: embed within `isa<ContractionOpInterface>` if possible / natural.
 bool isaContractionOpInterface(LinalgOp linalgOp);
 
 namespace detail {
 
 /// Verify that `op` conforms to ContractionOpInterface.
 LogicalResult verifyContractionInterface(Operation *op);

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

[{
       build($_builder, $_state, rangeType, min, max, step);
     }]>];
 
   // Fully specified by traits.
   let verifier = ?;
   let assemblyFormat = "$min `:` $max `:` $step attr-dict `:` type(results)";
 }
 
-class Linalg_ReshapeLikeOp<string mnemonic, list<OpTrait> traits = []> :
-    Linalg_Op<mnemonic, !listconcat(traits, [NoSideEffect])> {
-  let builders = [
-    // Builders for a contracting reshape whose result type is computed from
-    // `src` and `reassociation`.
-    OpBuilder<(ins "Value":$src,
-      "ArrayRef<ReassociationIndices>":$reassociation,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
-    OpBuilder<(ins "Value":$src,
-      "ArrayRef<ReassociationExprs>":$reassociation,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
-    [{
-      auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
-      build($_builder, $_state, src, reassociationMaps, attrs);
-    }]>,
-
-    // Builders for a reshape whose result type is passed explicitly. This may
-    // be either a contracting or expanding reshape.
-    OpBuilder<(ins "Type":$resultType, "Value":$src,
-      "ArrayRef<ReassociationIndices>":$reassociation,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
-    [{
-      build($_builder, $_state, resultType, src, attrs);
-      $_state.addAttribute("reassociation",
-                          getReassociationIndicesAttribute($_builder, reassociation));
-    }]>,
-    OpBuilder<(ins "Type":$resultType, "Value":$src,
-      "ArrayRef<ReassociationExprs>":$reassociation,
-      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
-    [{
-      auto reassociationMaps =
-          convertReassociationMapsToIndices($_builder, reassociation);
-      build($_builder, $_state, resultType, src, reassociationMaps, attrs);
-    }]>
-  ];
-
-  code commonExtraClassDeclaration = [{
-    static StringRef getReassociationAttrName() { return "reassociation"; }
-    SmallVector<AffineMap, 4> getReassociationMaps();
-    SmallVector<ReassociationExprs, 4> getReassociationExprs();
-    SmallVector<ReassociationIndices, 4> getReassociationIndices() {
-      SmallVector<ReassociationIndices, 4> reassociationIndices;
-      for (auto attr : reassociation())
-        reassociationIndices.push_back(llvm::to_vector<2>(
-            llvm::map_range(attr.cast<ArrayAttr>(), [&](Attribute indexAttr) {
-              return indexAttr.cast<IntegerAttr>().getInt();
-            })));
-      return reassociationIndices;
-    };
-  }];
-
-  let parser = [{ return ::parseReshapeLikeOp(parser, result); }];
-}
-
-def IndexListArrayAttr :
-  TypedArrayAttrBase<I64ArrayAttr, "Array of 64-bit integer array attributes">;
-
-class Linalg_TensorReshapeOp<string mnemonic> : Linalg_ReshapeLikeOp<
-    mnemonic,
-    [DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]>,
-    Arguments<(ins AnyTensor:$src,
-                   IndexListArrayAttr:$reassociation)>,
-    Results<(outs AnyTensor:$result)> {
-  let extraClassDeclaration = commonExtraClassDeclaration # [{
-    RankedTensorType getSrcType() {
-      return src().getType().cast<RankedTensorType>();
-    }
-    RankedTensorType getResultType() {
-      return result().getType().cast<RankedTensorType>();
-    }
-  }];
-  let hasFolder = 1;
-  let hasCanonicalizer = 1;
-  let printer = [{ return ::print(p, *this); }];
-  let parser = [{ return ::parseReshapeLikeOp(parser, result); }];
-}
-
-def Linalg_TensorExpandShapeOp : Linalg_TensorReshapeOp<"tensor_expand_shape"> {
-  let summary = "operation to produce a tensor with a higher rank";
-  let description = [{
-    The `linalg.tensor_expand_shape` op produces a new tensor with a higher
-    rank whose sizes are a reassociation of the original `src`.
-
-    A reassociation is defined as a continuous grouping of dimensions and is
-    represented with an array of I64ArrayAttr attribute.
-
-    The verification rule is that the reassociation maps are applied to the
-    result tensor with the higher rank to obtain the operand tensor with the
-    smaller rank.
-
-    The operand tensor type of a reshape can be zero-ranked if the result
-    tensor type is statically shaped with all dimensions being unit extent. In
-    such cases the reassociation map is empty.
-
-    Examples:
-
-    ```mlir
-    // Dimension expansion i -> (i', j') and (k) -> (k')
-    %b = linalg.tensor_expand_shape %a [[0, 1], [2]]
-        : tensor<?x?xf32> into tensor<?x?x?xf32>
-    ```
-  }];
-}
-
-def Linalg_TensorCollapseShapeOp : Linalg_TensorReshapeOp<"tensor_collapse_shape"> {
-  let summary = "operation to produce a tensor with a smaller rank";
-  let description = [{
-    The `linalg.tensor_collapse_shape` op produces a new tensor with a smaller
-    rank whose sizes are a reassociation of the original `src`.
-
-    A reassociation is defined as a continuous grouping of dimensions and is
-    represented with an array of I64ArrayAttr attribute.
-
-    The verification rule is that the reassociation maps are applied to the
-    operand tensor with the higher rank to obtain the result tensor with the
-    smaller rank.
-
-    The result tensor type of a reshape can be zero-ranked if the operand
-    tensor type is statically shaped with all dimensions being unit extent. In
-    such case the reassociation map is empty.
-
-    Examples:
-
-    ```mlir
-    // Dimension collapse (i, j) -> i' and k -> k'
-    %b = linalg.tensor_collapse_shape %a [[0, 1], [2]]
-        : tensor<?x?x?xf32> into tensor<?x?xf32>
-    ```
-  }];
-}
-
 def Linalg_YieldOp : Linalg_Op<"yield", [NoSideEffect, ReturnLike, Terminator]>,
     Arguments<(ins Variadic<AnyType>:$values)> {
   let summary = "Linalg yield operation";
   let description = [{
     `linalg.yield` is a special terminator operation for blocks inside regions
     in `linalg` generic ops. It returns values to the immediately enclosing
     `linalg` generic op.
 

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

def FillOp : LinalgStructured_Op<"fill", []> {
 
   let builders = [
     OpBuilder<(ins "Value":$value, "Value":$output)>
   ];
 
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Generic Linalg ops.
 //===----------------------------------------------------------------------===//
 
 def GenericOp : LinalgStructuredBase_Op<"generic", [AttrSizedOperandSegments]> {
   let description = [{

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let assemblyFormat = [{
     $source `(` $shape `)` attr-dict `:` functional-type(operands, results)
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // ExpandShapeOp / CollapseShapeOp
 //===----------------------------------------------------------------------===//
 
-def IndexListArrayAttr :
-  TypedArrayAttrBase<I64ArrayAttr, "Array of 64-bit integer array attributes">;
-
 class MemRef_ReassociativeReshapeOp<string mnemonic, list<OpTrait> traits = []> :
     MemRef_Op<mnemonic, !listconcat(traits,
       [NoSideEffect, ViewLikeOpInterface])>,
     Arguments<(ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation)>,
     Results<(outs AnyStridedMemRef:$result)>{
   let builders = [
     // Builders for a contracting reshape whose result type is computed from
     // `src` and `reassociation`.

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td

def ResolveShapedTypeResultDims : Pass<"resolve-shaped-type-result-dims"> {
   let description = [{
     The pass resolves memref.dim of result of operations that
     implement the `InferShapedTypeOpInterface` or
     `ReifyRankedShapedTypeOpInterface` in terms of shapes of its
     operands.
   }];
   let constructor = "mlir::memref::createResolveShapedTypeResultDimsPass()";
   let dependentDialects = [
-    "memref::MemRefDialect", "tensor::TensorDialect"
+    "AffineDialect", "memref::MemRefDialect", "tensor::TensorDialect"
   ];
 }
 
 #endif // MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES
 

mlir/include/mlir/Dialect/Tensor/IR/Tensor.h

 //===- Tensor.h - Tensor dialect --------------------------------*- C++ -*-===//
 //
 // 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_TENSOR_IR_TENSOR_H_
 #define MLIR_DIALECT_TENSOR_IR_TENSOR_H_
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"

mlir/include/mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h

+//===- TensorInferTypeOpInterfaceImpl.h - ---------------------------------===//
+//
+// 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 InferTypeOp interface for TensorOps with ExternalModel.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TENSOR_IR_TENSORINFERTYPEOPINTERFACEIMPL_H_
+#define MLIR_DIALECT_TENSOR_IR_TENSORINFERTYPEOPINTERFACEIMPL_H_
+
+#include "mlir/IR/Dialect.h"
+
+namespace mlir {
+namespace tensor {
+
+/// Registers external models for Infer Type interfaces for tensor ops.
+/// Currently, it registers:
+///
+/// * ReifyRankedShapedTypeOpInterface for `tensor.collapse_shape`.
+/// * ReifyRankedShapedTypeOpInterface for `tensor.expand_shape`.
+///
+/// Unfortunately, a "normal" internal registration is not possible at the
+/// moment, because of the dependency of the interface implementation for these
+/// ops on `affine.apply` and Affine dialect already depends on TensorOps. In
+/// order to break the cyclic dependency (TensorOps->AffineOps->TensorOps) the
+/// implementation is moved to a separate library.
+void registerInferTypeOpInterfaceExternalModels(
+    mlir::DialectRegistry &registry);
+
+} // namespace tensor
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TENSOR_IR_TENSORINFERTYPEOPINTERFACEIMPL_H_

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

def Tensor_ReshapeOp: Tensor_Op<"reshape", [NoSideEffect]>  {
   }];
 
   let assemblyFormat = [{
     $source `(` $shape `)` attr-dict `:` functional-type(operands, results)
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// ExpandShapeOp / CollapseShapeOp
+//===----------------------------------------------------------------------===//
+
+class Tensor_ReassociativeReshapeOp<string mnemonic, list<OpTrait> traits = []> :
+    Tensor_Op<mnemonic, !listconcat(traits, [NoSideEffect])>,
+    Arguments<(ins AnyTensor:$src, IndexListArrayAttr:$reassociation)>,
+    Results<(outs AnyTensor:$result)> {
+  let builders = [
+    // Builders for a contracting reshape whose result type is computed from
+    // `src` and `reassociation`.
+    OpBuilder<(ins "Value":$src,
+      "ArrayRef<ReassociationIndices>":$reassociation,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
+    OpBuilder<(ins "Value":$src,
+      "ArrayRef<ReassociationExprs>":$reassociation,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
+    [{
+      auto reassociationMaps =
+          convertReassociationMapsToIndices($_builder, reassociation);
+      build($_builder, $_state, src, reassociationMaps, attrs);
+    }]>,
+
+    // Builders for a reshape whose result type is passed explicitly. This may
+    // be either a contracting or expanding reshape.
+    OpBuilder<(ins "Type":$resultType, "Value":$src,
+      "ArrayRef<ReassociationIndices>":$reassociation,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
+    [{
+      build($_builder, $_state, resultType, src, attrs);
+      $_state.addAttribute("reassociation",
+          getReassociationIndicesAttribute($_builder, reassociation));
+    }]>,
+    OpBuilder<(ins "Type":$resultType, "Value":$src,
+      "ArrayRef<ReassociationExprs>":$reassociation,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
+    [{
+      auto reassociationMaps =
+          convertReassociationMapsToIndices($_builder, reassociation);
+      build($_builder, $_state, resultType, src, reassociationMaps, attrs);
+    }]>
+  ];
+
+  code commonExtraClassDeclaration = [{
+    static StringRef getReassociationAttrName() { return "reassociation"; }
+    SmallVector<AffineMap, 4> getReassociationMaps();
+    SmallVector<ReassociationExprs, 4> getReassociationExprs();
+    SmallVector<ReassociationIndices, 4> getReassociationIndices() {
+      SmallVector<ReassociationIndices, 4> reassociationIndices;
+      for (auto attr : reassociation())
+        reassociationIndices.push_back(llvm::to_vector<2>(
+            llvm::map_range(attr.cast<ArrayAttr>(), [&](Attribute indexAttr) {
+              return indexAttr.cast<IntegerAttr>().getInt();
+            })));
+      return reassociationIndices;
+    };
+    RankedTensorType getSrcType() {
+      return src().getType().cast<RankedTensorType>();
+    }
+    RankedTensorType getResultType() {
+      return result().getType().cast<RankedTensorType>();
+    }
+  }];
+
+  let hasFolder = 1;
+  let hasCanonicalizer = 1;
+  let printer = [{ return ::print(p, *this); }];
+  let parser = [{ return ::parseReshapeLikeOp(parser, result); }];
+}
+
+def Tensor_ExpandShapeOp : Tensor_ReassociativeReshapeOp<"expand_shape"> {
+  let summary = "operation to produce a tensor with a higher rank";
+  let description = [{
+    The `tensor.expand_shape` op produces a new tensor with a higher
+    rank whose sizes are a reassociation of the original `src`.
+
+    A reassociation is defined as a continuous grouping of dimensions and is
+    represented with an array of I64ArrayAttr attribute.
+
+    The verification rule is that the reassociation maps are applied to the
+    result tensor with the higher rank to obtain the operand tensor with the
+    smaller rank.
+
+    The operand tensor type of a reshape can be zero-ranked if the result
+    tensor type is statically shaped with all dimensions being unit extent. In
+    such cases the reassociation map is empty.
+
+    Examples:
+
+    ```mlir
+    // Dimension expansion i -> (i', j') and (k) -> (k')
+    %b = tensor.expand_shape %a [[0, 1], [2]]
+        : tensor<?x?xf32> into tensor<?x?x?xf32>
+    ```
+  }];
+  let extraClassDeclaration = commonExtraClassDeclaration;
+}
+
+def Tensor_CollapseShapeOp : Tensor_ReassociativeReshapeOp<"collapse_shape"> {
+  let summary = "operation to produce a tensor with a smaller rank";
+  let description = [{
+    The `tensor.collapse_shape` op produces a new tensor with a smaller
+    rank whose sizes are a reassociation of the original `src`.
+
+    A reassociation is defined as a continuous grouping of dimensions and is
+    represented with an array of I64ArrayAttr attribute.
+
+    The verification rule is that the reassociation maps are applied to the
+    operand tensor with the higher rank to obtain the result tensor with the
+    smaller rank.
+
+    The result tensor type of a reshape can be zero-ranked if the operand
+    tensor type is statically shaped with all dimensions being unit extent. In
+    such case the reassociation map is empty.
+
+    Examples:
+
+    ```mlir
+    // Dimension collapse (i, j) -> i' and k -> k'
+    %b = tensor.collapse_shape %a [[0, 1], [2]]
+        : tensor<?x?x?xf32> into tensor<?x?xf32>
+    ```
+  }];
+  let extraClassDeclaration = commonExtraClassDeclaration;
+}
+
+
+//===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_YieldOp : Tensor_Op<"yield",
     [NoSideEffect, ReturnLike, Terminator,
      HasParent<"::mlir::tensor::GenerateOp">]> {
   let summary = "Yield a value from a region";
   let description = [{

mlir/include/mlir/IR/OpBase.td

def F64ArrayAttr : TypedArrayAttrBase<F64Attr, "64-bit float array attribute"> {
   let constBuilderCall = "$_builder.getF64ArrayAttr($0)";
 }
 def StrArrayAttr : TypedArrayAttrBase<StrAttr, "string array attribute"> {
   let constBuilderCall = "$_builder.getStrArrayAttr($0)";
 }
 def TypeArrayAttr : TypedArrayAttrBase<TypeAttr, "type array attribute"> {
   let constBuilderCall = "$_builder.getTypeArrayAttr($0)";
 }
+def IndexListArrayAttr :
+  TypedArrayAttrBase<I64ArrayAttr, "Array of 64-bit integer array attributes">;
 
 // Attribute information for an Attribute field within a StructAttr.
 class StructFieldAttr<string thisName, Attr thisType> {
   // Name of this field in the StructAttr.
   string name = thisName;
 
   // Attribute type wrapped by the struct attr.
   Attr type = thisType;

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/Quant/QuantOps.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 #include "mlir/IR/Dialect.h"
 
 namespace mlir {
 
 /// Add all the MLIR dialects to the provided registry.
registry.insert<acc::OpenACCDialect,
                   NVVM::NVVMDialect,
                   ROCDL::ROCDLDialect,
                   shape::ShapeDialect,
                   sparse_tensor::SparseTensorDialect,
                   tensor::TensorDialect,
                   tosa::TosaDialect,
                   x86vector::X86VectorDialect>();
   // clang-format on
+  tensor::registerInferTypeOpInterfaceExternalModels(registry);
 }
 
 /// Append all the MLIR dialects to the registry contained in the given context.
 inline void registerAllDialects(MLIRContext &context) {
   DialectRegistry registry;
   registerAllDialects(registry);
   context.appendDialectRegistry(registry);
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLDIALECTS_H_

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

matchAndRewrite(tosa::ReshapeOp reshape, OpAdaptor adaptor,
 
     SmallVector<int64_t> intermediateShape;
     if (!findIntermediateShape(operandTy.getShape(), resultTy.getShape(),
                                intermediateShape, isDynamic)) {
       return rewriter.notifyMatchFailure(
           reshape, "tosa.reshape Cannot collapse into given shape");
     }
 
-    rewriter.replaceOpWithNewOp<linalg::TensorCollapseShapeOp>(
+    rewriter.replaceOpWithNewOp<tensor::CollapseShapeOp>(
         reshape, resultTy, adaptor.getOperands()[0], reassociationMap);
     return success();
   }
 };
 
 class ReshapeConverterExpand : public OpConversionPattern<tosa::ReshapeOp> {
 public:
   using OpConversionPattern<tosa::ReshapeOp>::OpConversionPattern;
matchAndRewrite(tosa::ReshapeOp reshape, OpAdaptor adaptor,
 
     SmallVector<int64_t> intermediateShape;
     if (!findIntermediateShape(operandTy.getShape(), resultTy.getShape(),
                                intermediateShape, isDynamic) ||
         intermediateShape != operandTy.getShape()) {
       return rewriter.notifyMatchFailure(
           reshape, "tosa.reshape Cannot expand into given shape");
     }
-    rewriter.replaceOpWithNewOp<linalg::TensorExpandShapeOp>(
+    rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
         reshape, resultTy, adaptor.getOperands()[0], reassociationMap);
     return success();
   }
 };
 
 class ReshapeConverterCollapseExpand
     : public OpConversionPattern<tosa::ReshapeOp> {
 public:

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 
 AffineApplyOp mlir::makeComposedAffineApply(OpBuilder &b, Location loc,
                                             AffineExpr e, ValueRange values) {
   return makeComposedAffineApply(
       b, loc, AffineMap::inferFromExprList(ArrayRef<AffineExpr>{e}).front(),
       values);
 }
 
+/// Fully compose map with operands and canonicalize the result.
+/// Return the `createOrFold`'ed AffineApply op.
+static Value createFoldedComposedAffineApply(OpBuilder &b, Location loc,
+                                             AffineMap map,
+                                             ValueRange operandsRef) {
+  SmallVector<Value, 4> operands(operandsRef.begin(), operandsRef.end());
+  fullyComposeAffineMapAndOperands(&map, &operands);
+  canonicalizeMapAndOperands(&map, &operands);
+  return b.createOrFold<AffineApplyOp>(loc, map, operands);
+}
+
+SmallVector<Value, 4> mlir::applyMapToValues(OpBuilder &b, Location loc,
+                                             AffineMap map, ValueRange values) {
+  SmallVector<Value, 4> res;
+  res.reserve(map.getNumResults());
+  unsigned numDims = map.getNumDims(), numSym = map.getNumSymbols();
+  // For each `expr` in `map`, applies the `expr` to the values extracted from
+  // ranges. If the resulting application can be folded into a Value, the
+  // folding occurs eagerly.
+  for (auto expr : map.getResults()) {
+    AffineMap map = AffineMap::get(numDims, numSym, expr);
+    res.push_back(createFoldedComposedAffineApply(b, loc, map, values));
+  }
+  return res;
+}
+
 // A symbol may appear as a dim in affine.apply operations. This function
 // canonicalizes dims that are valid symbols into actual symbols.
 template <class MapOrSet>
 static void canonicalizePromotedSymbols(MapOrSet *mapOrSet,
                                         SmallVectorImpl<Value> *operands) {
   if (!mapOrSet || operands->empty())
     return;
 

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

 OpOperandVector::operator SmallVector<Value>() {
   SmallVector<Value> result;
   result.reserve(this->size());
   llvm::transform(*this, std::back_inserter(result),
                   [](OpOperand *opOperand) { return opOperand->get(); });
   return result;
 }
 
-/// Fully compose map with operands and canonicalize the result.
-/// Return the `createOrFold`'ed AffineApply op.
-static Value createFoldedComposedAffineApply(OpBuilder &b, Location loc,
-                                             AffineMap map,
-                                             ValueRange operandsRef) {
-  SmallVector<Value, 4> operands(operandsRef.begin(), operandsRef.end());
-  fullyComposeAffineMapAndOperands(&map, &operands);
-  canonicalizeMapAndOperands(&map, &operands);
-  return b.createOrFold<AffineApplyOp>(loc, map, operands);
-}
-
-SmallVector<Value, 4> mlir::linalg::applyMapToValues(OpBuilder &b, Location loc,
-                                                     AffineMap map,
-                                                     ValueRange values) {
-  SmallVector<Value, 4> res;
-  res.reserve(map.getNumResults());
-  unsigned numDims = map.getNumDims(), numSym = map.getNumSymbols();
-  // For each `expr` in `map`, applies the `expr` to the values extracted from
-  // ranges. If the resulting application can be folded into a Value, the
-  // folding occurs eagerly.
-  for (auto expr : map.getResults()) {
-    AffineMap map = AffineMap::get(numDims, numSym, expr);
-    res.push_back(createFoldedComposedAffineApply(b, loc, map, values));
-  }
-  return res;
-}
-
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 static Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
                                int64_t dim) {
   if (source.getType().isa<UnrankedMemRefType, MemRefType>())
     return b.createOrFold<memref::DimOp>(loc, source, dim);
   if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
     return b.createOrFold<tensor::DimOp>(loc, source, dim);

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

static ParseResult parseNamedStructuredOp(OpAsmParser &parser,
                                           OperationState &result);
 
 static void printNamedStructuredOpResults(OpAsmPrinter &p,
                                           TypeRange resultTypes);
 
 template <typename NamedStructuredOpType>
 static void printNamedStructuredOp(OpAsmPrinter &p, NamedStructuredOpType op);
 
-/// Helper function to convert a vector of `OpFoldResult`s into a vector of
-/// `Value`s.
-static SmallVector<Value> getAsValues(OpBuilder &b, Location loc,
-                                      ArrayRef<OpFoldResult> valueOrAttrVec) {
-  return llvm::to_vector<4>(
-      llvm::map_range(valueOrAttrVec, [&](OpFoldResult value) -> Value {
-        return getValueOrCreateConstantIndexOp(b, loc, value);
-      }));
-}
-
 /// This is a common class used for patterns of the form
 /// ```
 ///    someop(memrefcast(%src)) -> someop(%src)
 /// ```
 /// It folds the source of the memref.cast into the root operation directly.
 static LogicalResult foldMemRefCast(Operation *op) {
   bool folded = false;
   for (OpOperand &operand : op->getOpOperands()) {

 void FillOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   if (output().getType().isa<MemRefType>())
     effects.emplace_back(MemoryEffects::Write::get(), output(),
                          SideEffects::DefaultResource::get());
 }
 
+namespace {
+
+/// Fold linalg.fill -> tensor.expand/collapse_shape chain.
+///
+/// For such op chains, we can create new linalg.fill ops with the result
+/// type of the tensor.expand/collapse_shape op.
+template <typename TensorReshapeOp>
+struct FoldFillWithTensorReshape : OpRewritePattern<TensorReshapeOp> {
+  using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
+                                PatternRewriter &rewriter) const override {
+    auto oldFill = reshapeOp.src().template getDefiningOp<FillOp>();
+    if (!oldFill)
+      return failure();
+
+    Location loc = oldFill.getLoc();
+    auto newInit = rewriter.create<TensorReshapeOp>(
+        loc, reshapeOp.getResultType(), oldFill.output(),
+        reshapeOp.reassociation());
+    rewriter.replaceOpWithNewOp<FillOp>(reshapeOp, oldFill.value(), newInit);
+
+    return success();
+  }
+};
+
+} // namespace
+
+void FillOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                         MLIRContext *context) {
+  results.add<FoldFillWithTensorReshape<tensor::CollapseShapeOp>,
+              FoldFillWithTensorReshape<tensor::ExpandShapeOp>>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // GenericOps
 //===----------------------------------------------------------------------===//
 void GenericOp::build(
     OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes,
     ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
     ArrayRef<StringRef> iteratorTypes, StringRef doc, StringRef libraryCall,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild,
struct FoldInitTensorWithTensorReshapeOp
   using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     if (!reshapeOp.src().template getDefiningOp<InitTensorOp>())
       return failure();
     Location loc = reshapeOp.getLoc();
     ReifiedRankedShapedTypeDims resultShapes;
-    if (failed(reshapeOp.reifyResultShapes(rewriter, resultShapes)) ||
+    ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
+        dyn_cast<ReifyRankedShapedTypeOpInterface>(reshapeOp.getOperation());
+    if (failed(reifyShapedTypeInterface.reifyResultShapes(rewriter,
+                                                          resultShapes)) ||
         !llvm::hasSingleElement(resultShapes))
       return failure();
     Value initTensor = rewriter.create<InitTensorOp>(
         loc, getAsOpFoldResult(resultShapes[0]),
         reshapeOp.getResultType().getElementType());
     if (initTensor.getType() != reshapeOp.getResultType()) {
       rewriter.replaceOpWithNewOp<tensor::CastOp>(
           reshapeOp, reshapeOp.getResultType(), initTensor);
LogicalResult matchAndRewrite(tensor::DimOp dimOp,
     return success();
   }
 };
 } // namespace
 
 void InitTensorOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
   results.add<FoldInitTensorWithDimOp, FoldInitTensorWithExtractSliceOp,
-              FoldInitTensorWithTensorReshapeOp<TensorExpandShapeOp>,
-              FoldInitTensorWithTensorReshapeOp<TensorCollapseShapeOp>,
+              FoldInitTensorWithTensorReshapeOp<tensor::ExpandShapeOp>,
+              FoldInitTensorWithTensorReshapeOp<tensor::CollapseShapeOp>,
               ReplaceStaticShapeDims>(context);
 }
 
 LogicalResult InitTensorOp::reifyResultShapes(
     OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   auto shapes = llvm::to_vector<4>(llvm::map_range(
       llvm::seq<int64_t>(0, getType().getRank()), [&](int64_t dim) -> Value {
         if (isDynamicSize(dim))

 OpFoldResult PadTensorOp::fold(ArrayRef<Attribute>) {
   if (getResultType().hasStaticShape() && getResultType() == getSourceType() &&
       !nofold())
     return source();
   return {};
 }
 
 //===----------------------------------------------------------------------===//
-// ReshapeOp
-//===----------------------------------------------------------------------===//
-
-static void print(OpAsmPrinter &p, linalg::TensorExpandShapeOp op) {
-  ::mlir::printReshapeOp<linalg::TensorExpandShapeOp>(p, op);
-}
-
-static void print(OpAsmPrinter &p, linalg::TensorCollapseShapeOp op) {
-  ::mlir::printReshapeOp<linalg::TensorCollapseShapeOp>(p, op);
-}
-
-template <typename AffineExprTy>
-unsigned getMaxPosOfType(ArrayRef<ReassociationExprs> exprArrays) {
-  unsigned pos = 0;
-  for (const auto &exprs : exprArrays) {
-    for (auto expr : exprs) {
-      expr.walk([&pos](AffineExpr e) {
-        if (auto d = e.dyn_cast<AffineExprTy>())
-          pos = std::max(pos, d.getPosition());
-      });
-    }
-  }
-  return pos;
-}
-
-SmallVector<AffineMap, 4> TensorCollapseShapeOp::getReassociationMaps() {
-  return getSymbolLessAffineMaps(getReassociationExprs());
-}
-SmallVector<ReassociationExprs, 4>
-TensorCollapseShapeOp::getReassociationExprs() {
-  return convertReassociationIndicesToExprs(getContext(),
-                                            getReassociationIndices());
-}
-SmallVector<AffineMap, 4> TensorExpandShapeOp::getReassociationMaps() {
-  return getSymbolLessAffineMaps(getReassociationExprs());
-}
-SmallVector<ReassociationExprs, 4>
-TensorExpandShapeOp::getReassociationExprs() {
-  return convertReassociationIndicesToExprs(getContext(),
-                                            getReassociationIndices());
-}
-
-/// For reshape op compute the shape at dimension `dimIndex` of the output in
-/// terms of shape of the `src`, when the reshape op is a collapsing
-/// operation. It is the product of the shape of the collapsed dimensions of the
-/// `src`.
-static OpFoldResult
-getCollapsedOutputDimFromInputShape(OpBuilder &builder, Location loc,
-                                    int64_t dimIndex, Value src,
-                                    ArrayRef<AffineMap> reassociationMap) {
-  AffineMap map = reassociationMap[dimIndex];
-  unsigned startPos =
-      map.getResults().front().cast<AffineDimExpr>().getPosition();
-  unsigned endPos = map.getResults().back().cast<AffineDimExpr>().getPosition();
-  AffineExpr expr;
-  SmallVector<Value, 2> dynamicDims;
-  for (auto dim : llvm::seq_inclusive(startPos, endPos)) {
-    dynamicDims.push_back(builder.createOrFold<tensor::DimOp>(loc, src, dim));
-    AffineExpr currExpr = builder.getAffineSymbolExpr(dim - startPos);
-    expr = (expr ? expr * currExpr : currExpr);
-  }
-  return applyMapToValues(builder, loc,
-                          AffineMap::get(0, endPos - startPos + 1, expr),
-                          dynamicDims)[0];
-}
-
-/// Given the `src` of a collapsing reshape op and its reassociation maps,
-/// compute the shape of the result of the reshape.
-static SmallVector<OpFoldResult, 4> getCollapsedOutputShapeFromInputShape(
-    OpBuilder &builder, Location loc, Value src,
-    ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
-  return llvm::to_vector<4>(llvm::map_range(
-      llvm::seq<int64_t>(0, dstStaticShape.size()), [&](int64_t dim) {
-        return getCollapsedOutputDimFromInputShape(builder, loc, dim, src,
-                                                   reassociation);
-      }));
-}
-
-/// Compute a map that for a given dimension of the expanded type gives the
-/// dimension in the collapsed type it maps to. Essentially its the inverse of
-/// the `reassocation` maps.
-static llvm::DenseMap<int64_t, int64_t>
-getExpandedDimToCollapsedDimMap(ArrayRef<AffineMap> reassociation) {
-  llvm::DenseMap<int64_t, int64_t> expandedDimToCollapsedDim;
-  for (auto map : enumerate(reassociation)) {
-    unsigned startPos =
-        map.value().getResults().front().cast<AffineDimExpr>().getPosition();
-    unsigned endPos =
-        map.value().getResults().back().cast<AffineDimExpr>().getPosition();
-    for (auto dim : llvm::seq_inclusive(startPos, endPos)) {
-      expandedDimToCollapsedDim[dim] = map.index();
-    }
-  }
-  return expandedDimToCollapsedDim;
-}
-
-/// For an expanding reshape op, compute the value for a dimension of the output
-/// from the shape of the input.
-static OpFoldResult getExpandedOutputDimFromInputShape(
-    OpBuilder &builder, Location loc, int64_t dimIndex, Value src,
-    ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation,
-    llvm::DenseMap<int64_t, int64_t> &expandedDimToCollapsedDim) {
-  if (!ShapedType::isDynamic(dstStaticShape[dimIndex])) {
-    return builder.getI64IntegerAttr(dstStaticShape[dimIndex]);
-  }
-  unsigned sourceDimPos = expandedDimToCollapsedDim[dimIndex];
-  unsigned startPos = reassociation[sourceDimPos]
-                          .getResults()
-                          .front()
-                          .cast<AffineDimExpr>()
-                          .getPosition();
-  unsigned endPos = reassociation[sourceDimPos]
-                        .getResults()
-                        .back()
-                        .cast<AffineDimExpr>()
-                        .getPosition();
-  int64_t linearizedStaticDim = 1;
-  for (auto d :
-       llvm::enumerate(dstStaticShape.slice(startPos, endPos - startPos + 1))) {
-    if (d.index() + startPos == static_cast<unsigned>(dimIndex))
-      continue;
-    assert(!ShapedType::isDynamic(d.value()) &&
-           "single dimension cannot be expanded into multiple dynamic "
-           "dimensions");
-    linearizedStaticDim *= d.value();
-  }
-  Value sourceDim = builder.create<tensor::DimOp>(loc, src, sourceDimPos);
-  return applyMapToValues(
-      builder, loc,
-      AffineMap::get(
-          0, 1, builder.getAffineSymbolExpr(0).floorDiv(linearizedStaticDim)),
-      sourceDim)[0];
-}
-
-/// Given the `src` of an expanding reshape op, the reassociation maps and the
-/// result type, compute the shape of the result of the reshape.
-static SmallVector<OpFoldResult, 4> getExpandedOutputShapeFromInputShape(
-    OpBuilder &builder, Location loc, Value src,
-    ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
-  llvm::DenseMap<int64_t, int64_t> expandedDimToCollapsedDim =
-      getExpandedDimToCollapsedDimMap(reassociation);
-  return llvm::to_vector<4>(llvm::map_range(
-      llvm::seq<int64_t>(0, dstStaticShape.size()), [&](int64_t dim) {
-        return getExpandedOutputDimFromInputShape(builder, loc, dim, src,
-                                                  dstStaticShape, reassociation,
-                                                  expandedDimToCollapsedDim);
-      }));
-}
-
-static SmallVector<OpFoldResult, 4>
-getReshapeOutputShapeFromInputShape(OpBuilder &builder, Location loc, Value src,
-                                    ArrayRef<int64_t> dstStaticShape,
-                                    ArrayRef<AffineMap> reassocation) {
-  return dstStaticShape.size() >
-                 static_cast<size_t>(src.getType().cast<ShapedType>().getRank())
-             ? getExpandedOutputShapeFromInputShape(
-                   builder, loc, src, dstStaticShape, reassocation)
-             : getCollapsedOutputShapeFromInputShape(
-                   builder, loc, src, dstStaticShape, reassocation);
-}
-
-//===----------------------------------------------------------------------===//
-// TensorReshapeOp
-//===----------------------------------------------------------------------===//
-
-/// Compute the RankedTensorType obtained by applying `reassociation` to `type`.
-static RankedTensorType
-computeTensorReshapeCollapsedType(RankedTensorType type,
-                                  ArrayRef<AffineMap> reassociation) {
-  auto shape = type.getShape();
-  SmallVector<int64_t, 4> newShape;
-  newShape.reserve(reassociation.size());
-
-  // Use the fact that reassociation is valid to simplify the logic: only use
-  // each map's rank.
-  assert(isReassociationValid(reassociation) && "invalid reassociation");
-  unsigned currentDim = 0;
-  for (AffineMap m : reassociation) {
-    unsigned dim = m.getNumResults();
-    auto band = shape.slice(currentDim, dim);
-    int64_t size = 1;
-    if (llvm::is_contained(band, ShapedType::kDynamicSize))
-      size = ShapedType::kDynamicSize;
-    else
-      for (unsigned d = 0; d < dim; ++d)
-        size *= shape[currentDim + d];
-    newShape.push_back(size);
-    currentDim += dim;
-  }
-
-  return RankedTensorType::get(newShape, type.getElementType());
-}
-
-void mlir::linalg::TensorCollapseShapeOp::build(
-    OpBuilder &b, OperationState &result, Value src,
-    ArrayRef<ReassociationIndices> reassociation,
-    ArrayRef<NamedAttribute> attrs) {
-  auto resultType = computeTensorReshapeCollapsedType(
-      src.getType().cast<RankedTensorType>(),
-      getSymbolLessAffineMaps(
-          convertReassociationIndicesToExprs(b.getContext(), reassociation)));
-  build(b, result, resultType, src, attrs);
-  result.addAttribute(getReassociationAttrName(),
-                      getReassociationIndicesAttribute(b, reassociation));
-}
-
-void mlir::linalg::TensorExpandShapeOp::build(
-    OpBuilder &b, OperationState &result, Value src,
-    ArrayRef<ReassociationIndices> reassociation,
-    ArrayRef<NamedAttribute> attrs) {
-  auto resultType = computeTensorReshapeCollapsedType(
-      src.getType().cast<RankedTensorType>(),
-      getSymbolLessAffineMaps(
-          convertReassociationIndicesToExprs(b.getContext(), reassociation)));
-  build(b, result, resultType, src, attrs);
-  result.addAttribute(getReassociationAttrName(),
-                      getReassociationIndicesAttribute(b, reassociation));
-}
-
-template <typename TensorReshapeOp,
-          bool isExpansion =
-              std::is_same<TensorReshapeOp, TensorExpandShapeOp>::value>
-static LogicalResult verifyTensorReshapeOp(TensorReshapeOp op,
-                                           RankedTensorType expandedType,
-                                           RankedTensorType collapsedType) {
-  if (failed(
-          verifyReshapeLikeTypes(op, expandedType, collapsedType, isExpansion)))
-    return failure();
-
-  auto maps = op.getReassociationMaps();
-  RankedTensorType expectedType =
-      computeTensorReshapeCollapsedType(expandedType, maps);
-  if (collapsedType != expectedType)
-    return op.emitOpError("expected collapsed type to be ")
-           << expectedType << ", but got " << collapsedType;
-  return success();
-}
-
-static LogicalResult verify(TensorExpandShapeOp op) {
-  return verifyTensorReshapeOp(op, op.getResultType(), op.getSrcType());
-}
-
-static LogicalResult verify(TensorCollapseShapeOp op) {
-  return verifyTensorReshapeOp(op, op.getSrcType(), op.getResultType());
-}
-
-namespace {
-/// Reshape of a splat constant can be replaced with a constant of the result
-/// type.
-template <typename TensorReshapeOp>
-struct FoldReshapeWithConstant : OpRewritePattern<TensorReshapeOp> {
-  using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
-                                PatternRewriter &rewriter) const override {
-    DenseElementsAttr attr;
-    if (!matchPattern(reshapeOp.src(), m_Constant(&attr)))
-      return failure();
-    if (!attr || !attr.isSplat())
-      return failure();
-    DenseElementsAttr newAttr = DenseElementsAttr::getFromRawBuffer(
-        reshapeOp.getResultType(), attr.getRawData(), true);
-    rewriter.replaceOpWithNewOp<arith::ConstantOp>(reshapeOp, newAttr);
-    return success();
-  }
-};
-
-/// Fold linalg.fill -> linalg.tensor_reshape chain.
-///
-/// For such op chains, we can create new linalg.fill ops with the result
-/// type of the linalg.tensor_reshape op.
-template <typename TensorReshapeOp>
-struct FoldFillWithTensorReshape : OpRewritePattern<TensorReshapeOp> {
-  using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
-                                PatternRewriter &rewriter) const override {
-    auto oldFill = reshapeOp.src().template getDefiningOp<FillOp>();
-    if (!oldFill)
-      return failure();
-
-    Location loc = oldFill.getLoc();
-    auto newInit = rewriter.create<TensorReshapeOp>(
-        loc, reshapeOp.getResultType(), oldFill.output(),
-        reshapeOp.reassociation());
-    rewriter.replaceOpWithNewOp<FillOp>(reshapeOp, oldFill.value(), newInit);
-
-    return success();
-  }
-};
-} // namespace
-
-void TensorExpandShapeOp::getCanonicalizationPatterns(
-    RewritePatternSet &results, MLIRContext *context) {
-  results
-      .add<CollapseReshapeOps<TensorExpandShapeOp>,
-           CollapseMixedReshapeOps<TensorExpandShapeOp, TensorCollapseShapeOp>,
-           FoldFillWithTensorReshape<TensorExpandShapeOp>,
-           FoldInitTensorWithTensorReshapeOp<TensorExpandShapeOp>,
-           FoldReshapeWithConstant<TensorExpandShapeOp>>(context);
-}
-
-void TensorCollapseShapeOp::getCanonicalizationPatterns(
-    RewritePatternSet &results, MLIRContext *context) {
-  results
-      .add<CollapseReshapeOps<TensorCollapseShapeOp>,
-           CollapseMixedReshapeOps<TensorCollapseShapeOp, TensorExpandShapeOp>,
-           FoldFillWithTensorReshape<TensorCollapseShapeOp>,
-           FoldInitTensorWithTensorReshapeOp<TensorCollapseShapeOp>,
-           FoldReshapeWithConstant<TensorCollapseShapeOp>>(context);
-}
-
-LogicalResult TensorExpandShapeOp::reifyResultShapes(
-    OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
-  auto resultShape =
-      getAsValues(b, getLoc(),
-                  getReshapeOutputShapeFromInputShape(
-                      b, getLoc(), src(), getResultType().getShape(),
-                      getReassociationMaps()));
-  reifiedReturnShapes.emplace_back(std::move(resultShape));
-  return success();
-}
-
-LogicalResult TensorCollapseShapeOp::reifyResultShapes(
-    OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
-  auto resultShape =
-      getAsValues(b, getLoc(),
-                  getReshapeOutputShapeFromInputShape(
-                      b, getLoc(), src(), getResultType().getShape(),
-                      getReassociationMaps()));
-  reifiedReturnShapes.emplace_back(std::move(resultShape));
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, linalg::YieldOp op) {
   if (op.getNumOperands() > 0)
     p << ' ' << op.getOperands();
   p.printOptionalAttrDict(op->getAttrs());
   if (op.getNumOperands() > 0)
std::string mlir::linalg::generateLibraryCallName(Operation *op) {
   ss << "_";
   auto types = op->getOperandTypes();
   llvm::interleave(
       types.begin(), types.end(), [&](Type t) { appendMangledType(ss, t); },
       [&]() { ss << "_"; });
   return ss.str();
 }
 
-// TODO: Consider making all this boilerplate easy to autogenerate
-// with Tablegen. This seems a desirable property in the context of
-// OpInterfaces where a Linalg "named" op **isa** LinalgOp.
-OpFoldResult TensorExpandShapeOp::fold(ArrayRef<Attribute> operands) {
-  return foldReshapeOp<TensorExpandShapeOp, TensorCollapseShapeOp>(*this,
-                                                                   operands);
-}
-OpFoldResult TensorCollapseShapeOp::fold(ArrayRef<Attribute> operands) {
-  return foldReshapeOp<TensorCollapseShapeOp, TensorExpandShapeOp>(*this,
-                                                                   operands);
-}
-
 //===----------------------------------------------------------------------===//
 // Support for named Linalg ops defined in ods-gen.
 //===----------------------------------------------------------------------===//
 
 /// Generic entry point to create the block for the region of a LinalgOp.
 /// This is used by both named structured ops created by ods-gen and by manually
 /// defined C++ ops.
 /// This is used by both builders and parsers.
if (kernelTy.getDimSize(3) != 1)
     return failure();
 
   // Collapse kernel dims.
   SmallVector<ReassociationIndices, 4> collapsedKernelDims = {
       getIndicesVector(0, 1), getIndicesVector(1, 2), getIndicesVector(2, 4)};
   auto newKernelTy = RankedTensorType::get(
       {kernelTy.getDimSize(0), kernelTy.getDimSize(1), kernelTy.getDimSize(2)},
       kernelTy.getElementType());
-  auto collapsedKernel = rewriter.create<linalg::TensorCollapseShapeOp>(
+  auto collapsedKernel = rewriter.create<tensor::CollapseShapeOp>(
       loc, newKernelTy, kernel, collapsedKernelDims);
 
   // Collapse init dims.
   SmallVector<ReassociationIndices, 4> collapsedInitDims = {
       getIndicesVector(0, 1), getIndicesVector(1, 2), getIndicesVector(2, 3),
       getIndicesVector(3, 5)};
   auto newInitTy =
       RankedTensorType::get({initTy.getDimSize(0), initTy.getDimSize(1),
                              initTy.getDimSize(2), initTy.getDimSize(3)},
                             initTy.getElementType());
-  auto collapsedInit = rewriter.create<linalg::TensorCollapseShapeOp>(
+  auto collapsedInit = rewriter.create<tensor::CollapseShapeOp>(
       loc, newInitTy, init, collapsedInitDims);
 
   Value newConv;
   if (isa<DepthwiseConv2DNhwcHwcmOp>(operation)) {
     newConv = rewriter
                   .create<DepthwiseConv2DNhwcHwcOp>(
                       loc, newInitTy, ValueRange{input, collapsedKernel},
                       ValueRange{collapsedInit}, stride, dilation)
                   .getResult(0);
   } else if (isa<DepthwiseConv2DNhwcHwcmQOp>(operation)) {
     newConv =
         rewriter
             .create<DepthwiseConv2DNhwcHwcQOp>(
                 loc, newInitTy, ValueRange{input, collapsedKernel, iZp, kZp},
                 ValueRange{collapsedInit}, stride, dilation)
             .getResult(0);
   }
 
   if (!newConv)
     return failure();
 
   // Expand dimensions back out to
-  rewriter.replaceOpWithNewOp<linalg::TensorExpandShapeOp>(
+  rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
       operation, resultTy, newConv, collapsedInitDims);
   return success();
 }
 
 struct SimplifyDepthwiseConvOp
     : public OpRewritePattern<DepthwiseConv2DNhwcHwcmOp> {
   using OpRewritePattern<DepthwiseConv2DNhwcHwcmOp>::OpRewritePattern;
 

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

 /// Conversion pattern that replaces `linalg.tensor_reshape` with
 /// `linalg.reshape`.
 template <typename TensorReshapeOp,
           typename Adaptor = typename TensorReshapeOp::Adaptor>
 class BufferizeTensorReshapeOp : public OpConversionPattern<TensorReshapeOp> {
 public:
   using OpConversionPattern<TensorReshapeOp>::OpConversionPattern;
   using ReshapeOp = typename std::conditional_t<
-      std::is_same<TensorReshapeOp, TensorExpandShapeOp>::value,
+      std::is_same<TensorReshapeOp, tensor::ExpandShapeOp>::value,
       memref::ExpandShapeOp, memref::CollapseShapeOp>;
 
   LogicalResult
   matchAndRewrite(TensorReshapeOp op, Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
     rewriter.replaceOpWithNewOp<ReshapeOp>(op,
                                            this->getTypeConverter()
                                                ->convertType(op.getType())
void runOnOperation() override {
     MLIRContext &context = getContext();
     ConversionTarget target(context);
     bufferization::BufferizeTypeConverter typeConverter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<arith::ArithmeticDialect, AffineDialect,
                            memref::MemRefDialect, StandardOpsDialect,
                            tensor::TensorDialect>();
-    target.addIllegalOp<InitTensorOp, tensor::ExtractSliceOp,
-                        tensor::InsertSliceOp, PadTensorOp>();
+    target.addIllegalOp<InitTensorOp, PadTensorOp, tensor::CollapseShapeOp,
+                        tensor::ExpandShapeOp, tensor::ExtractSliceOp,
+                        tensor::InsertSliceOp>();
 
     // Mark all Linalg operations illegal as long as they work on tensors.
     auto isLegalOperation = [&](Operation *op) {
       return typeConverter.isLegal(op);
     };
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(isLegalOperation);
     target
         .addDynamicallyLegalOp<vector::TransferReadOp, vector::TransferWriteOp>(
void mlir::linalg::populateLinalgBufferizePatterns(
     bufferization::BufferizeTypeConverter &typeConverter,
     RewritePatternSet &patterns) {
   // TODO: Drop this once tensor constants work in standard.
   // clang-format off
   patterns.add<
       BufferizeAnyLinalgOp,
       BufferizeFillOp,
       BufferizeInitTensorOp,
-      BufferizeTensorReshapeOp<TensorExpandShapeOp>,
-      BufferizeTensorReshapeOp<TensorCollapseShapeOp>,
+      BufferizeTensorReshapeOp<tensor::ExpandShapeOp>,
+      BufferizeTensorReshapeOp<tensor::CollapseShapeOp>,
       ExtractSliceOpConverter,
       InsertSliceOpConverter,
       VectorTransferReadOpConverter,
       VectorTransferWriteOpConverter
     >(typeConverter, patterns.getContext());
   // clang-format on
   patterns.add<GeneralizePadTensorOpPattern>(patterns.getContext());
 }

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

static Value sourceMaterializationCallback(OpBuilder &builder, Type type,
 
   // A detensored value is converted back by creating a new tensor from its
   // element(s).
   auto createNewTensorOp = builder.create<tensor::FromElementsOp>(
       loc, inputs[0].getType(), inputs[0]);
 
   // FromElementsOp results in a tensor<1xdtype>, we need to reshape that to
   // a tensor<dtype> instead.
-  return builder.create<linalg::TensorCollapseShapeOp>(
+  return builder.create<tensor::CollapseShapeOp>(
       loc, type, createNewTensorOp, ArrayRef<ReassociationExprs>{});
 }
 
 namespace {
 /// Defines the criteria a TensorType must follow in order to be considered
 /// "detensorable".
 ///
 /// NOTE: For now, only 0-D tensors are supported.
struct ExtractFromReshapeFromElements
   using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tensor::ExtractOp extract,
                                 PatternRewriter &rewriter) const final {
     if (!extract.indices().empty())
       return failure();
 
     auto tensorReshape =
-        extract.tensor().getDefiningOp<TensorCollapseShapeOp>();
+        extract.tensor().getDefiningOp<tensor::CollapseShapeOp>();
     if (tensorReshape == nullptr)
       return failure();
 
     auto tensorFromElements =
         tensorReshape.getOperand()
             .getDefiningOp<mlir::tensor::FromElementsOp>();
     if (tensorFromElements == nullptr)
       return failure();

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

struct ReplaceUnitExtents : public OpRewritePattern<GenericOp> {
   // Return the original value if the type is unchanged, or reshape it. Return a
   // nullptr if this is an unsupported type.
   Value maybeExpand(Value result, Type origResultType,
                     ArrayAttr reassociationMap, Location loc,
                     PatternRewriter &rewriter) const {
     if (origResultType == result.getType())
       return result;
     if (origResultType.isa<RankedTensorType>()) {
-      return rewriter.create<linalg::TensorExpandShapeOp>(
+      return rewriter.create<tensor::ExpandShapeOp>(
           loc, origResultType, result,
           convertAffineMapArrayToExprs(reassociationMap));
     }
     if (origResultType.isa<MemRefType>()) {
       return rewriter.create<memref::ExpandShapeOp>(
           loc, origResultType, result,
           convertAffineMapArrayToExprs(reassociationMap));
     }
Value maybeCollapse(Value operand, Type newInputOutputType,
     if (operandType == newInputOutputType)
       return operand;
     if (operandType.isa<MemRefType>()) {
       return rewriter.create<memref::CollapseShapeOp>(
           loc, newInputOutputType, operand,
           convertAffineMapArrayToExprs(reassociationMap));
     }
     if (operandType.isa<RankedTensorType>()) {
-      return rewriter.create<linalg::TensorCollapseShapeOp>(
+      return rewriter.create<tensor::CollapseShapeOp>(
           loc, newInputOutputType, operand,
           convertAffineMapArrayToExprs(reassociationMap));
     }
     return nullptr;
   };
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
     auto rankReducedType = tensor::ExtractSliceOp::inferRankReducedResultType(
                                reassociation->size(), sliceOp.getSourceType(),
                                offsets, sizes, strides)
                                .cast<RankedTensorType>();
 
     Location loc = sliceOp.getLoc();
     Value newSlice = rewriter.create<tensor::ExtractSliceOp>(
         loc, rankReducedType, sliceOp.source(), offsets, sizes, strides);
-    rewriter.replaceOpWithNewOp<TensorExpandShapeOp>(sliceOp, resultType,
-                                                     newSlice, *reassociation);
+    rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
+        sliceOp, resultType, newSlice, *reassociation);
     return success();
   }
 };
 
 /// Convert `insert_slice` operations to rank-reduced versions.
 struct UseRankReducedInsertSliceOp
     : public OpRewritePattern<tensor::InsertSliceOp> {
   using OpRewritePattern<tensor::InsertSliceOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tensor::InsertSliceOp insertOp,
                                 PatternRewriter &rewriter) const override {
     RankedTensorType sourceType = insertOp.getSourceType();
     SmallVector<OpFoldResult> offsets = insertOp.getMixedOffsets();
     SmallVector<OpFoldResult> sizes = insertOp.getMixedSizes();
     SmallVector<OpFoldResult> strides = insertOp.getMixedStrides();
     auto reassociation = getReassociationMapForFoldingUnitDims(sizes);
     if (!reassociation ||
         reassociation->size() == static_cast<size_t>(sourceType.getRank()))
       return failure();
     Location loc = insertOp.getLoc();
-    auto reshapedSource = rewriter.create<TensorCollapseShapeOp>(
+    auto reshapedSource = rewriter.create<tensor::CollapseShapeOp>(
         loc, insertOp.source(), *reassociation);
     rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
         insertOp, reshapedSource, insertOp.dest(), insertOp.getMixedOffsets(),
         insertOp.getMixedSizes(), insertOp.getMixedStrides());
     return success();
   }
 };
 } // namespace
 
 /// Patterns that are used to canonicalize the use of unit-extent dims for
 /// broadcasting.
 void mlir::linalg::populateFoldUnitExtentDimsPatterns(
     RewritePatternSet &patterns) {
   auto *context = patterns.getContext();
   patterns.add<FoldUnitDimLoops, ReplaceUnitExtents,
                UseRankReducedExtractSliceOp, UseRankReducedInsertSliceOp>(
       context);
-  TensorCollapseShapeOp::getCanonicalizationPatterns(patterns, context);
-  TensorExpandShapeOp::getCanonicalizationPatterns(patterns, context);
+  linalg::FillOp::getCanonicalizationPatterns(patterns, context);
+  linalg::InitTensorOp::getCanonicalizationPatterns(patterns, context);
+  tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
+  tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
 }
 
 namespace {
 /// Pass that removes unit-extent dims within generic ops.
 struct LinalgFoldUnitExtentDimsPass
     : public LinalgFoldUnitExtentDimsBase<LinalgFoldUnitExtentDimsPass> {
   void runOnFunction() override {
     FuncOp funcOp = getFunction();

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

 /// would be rewritten into:
 /// %0 = op ... : tensor<?x?x4x5xf32>
 /// with output index_map
 ///   `affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>`
 template <typename TensorReshapeOp>
 static AffineMap linearizeCollapsedDims(AffineMap sourceMap,
                                         TensorReshapeOp reshapeOp) {
   constexpr bool isExpanding =
-      std::is_same<TensorReshapeOp, TensorExpandShapeOp>::value;
+      std::is_same<TensorReshapeOp, tensor::ExpandShapeOp>::value;
   ArrayRef<int64_t> sourceShape =
       (isExpanding ? reshapeOp.getResultType().getShape()
                    : reshapeOp.getSrcType().getShape());
   SmallVector<AffineExpr> resultExprs;
   ArrayRef<AffineExpr> sourceExprs = sourceMap.getResults();
   MLIRContext *context = sourceMap.getContext();
 
   // Compute the result exprs based on the reassociation maps.
static AffineMap linearizeCollapsedDims(AffineMap sourceMap,
   int64_t maxDim = -1;
   int64_t maxSym = -1;
   getMaxDimAndSymbol<SmallVector<AffineExpr>>({resultExprs}, maxDim, maxSym);
   unsigned numDims = std::max(unsigned(maxDim + 1), sourceMap.getNumDims());
   unsigned numSyms = std::max(unsigned(maxSym + 1), sourceMap.getNumSymbols());
   return AffineMap::get(numDims, numSyms, resultExprs, context);
 }
 
-// TensorExpandShapeOp is fusable with its consumer (i.e. reshape as a
+// tensor::ExpandShapeOp is fusable with its consumer (i.e. reshape as a
 // producer). Fusing when operand has higher rank will require use of mods and
 // divs in the indexing maps of the fused op which would make it non-invertible.
 static bool isTensorReshapeOpFoldableByLinearization(
-    TensorExpandShapeOp expandOp, AffineMap useIndexMap, bool asProducer) {
+    tensor::ExpandShapeOp expandOp, AffineMap useIndexMap, bool asProducer) {
   if (!asProducer)
     return false;
   return useIndexMap.isPermutation();
 }
 
-// TensorCollapseShapeOp is fusable with its producer (i.e. reshape as a
+// tensor::CollapseShapeOp is fusable with its producer (i.e. reshape as a
 // consumer).
-static bool isTensorReshapeOpFoldableByLinearization(
-    TensorCollapseShapeOp collapseOp, AffineMap useIndexMap, bool asProducer) {
+static bool
+isTensorReshapeOpFoldableByLinearization(tensor::CollapseShapeOp collapseOp,
+                                         AffineMap useIndexMap,
+                                         bool asProducer) {
   if (asProducer)
     return false;
   return useIndexMap.isPermutation();
 }
 
 /// Check if the reshape operation is only expansion into/collapsing of
 /// unit-dimension.
 template <typename TensorReshapeOp>
 static bool isUnitDimExpansionOnly(TensorReshapeOp reshapeOp) {
   constexpr bool isExpanding =
-      std::is_same<TensorReshapeOp, TensorExpandShapeOp>::value;
+      std::is_same<TensorReshapeOp, tensor::ExpandShapeOp>::value;
   ArrayRef<int64_t> expandedShape =
       (isExpanding ? reshapeOp.getResultType().getShape()
                    : reshapeOp.getSrcType().getShape());
   for (auto &indices : reshapeOp.getReassociationIndices()) {
     unsigned numUnitDims = 0;
     for (int64_t position : indices)
       if (expandedShape[position] == 1)
         numUnitDims++;

 /// that those conditions have been satisfied.
 static Optional<SmallVector<Value>>
 fuseWithReshapeByExpansion(GenericOp genericOp, Operation *reshapeOp,
                            OpOperand *fusableOpOperand,
                            PatternRewriter &rewriter) {
   assert(isFusableWithReshapeByDimExpansion(genericOp, fusableOpOperand) &&
          "preconditions for fuse operation failed");
   // Check if reshape is expanding or collapsing.
-  auto expandingReshapeOp = dyn_cast<TensorExpandShapeOp>(*reshapeOp);
-  auto collapsingReshapeOp = dyn_cast<TensorCollapseShapeOp>(*reshapeOp);
+  auto expandingReshapeOp = dyn_cast<tensor::ExpandShapeOp>(*reshapeOp);
+  auto collapsingReshapeOp = dyn_cast<tensor::CollapseShapeOp>(*reshapeOp);
   bool isExpanding = (expandingReshapeOp != nullptr);
   RankedTensorType expandedType = isExpanding
                                       ? expandingReshapeOp.getResultType()
                                       : collapsingReshapeOp.getSrcType();
 
   ExpansionInfo expansionInfo;
   if (failed(expansionInfo.compute(
           genericOp, fusableOpOperand,
if (genericOp.isInputTensor(opOperand)) {
       AffineMap indexingMap = genericOp.getTiedIndexingMap(opOperand);
       RankedTensorType expandedOperandType =
           getExpandedType(opOperand->get().getType().cast<RankedTensorType>(),
                           indexingMap, expansionInfo);
       if (expandedOperandType != opOperand->get().getType()) {
         // Reshape the operand to get the right type.
         SmallVector<ReassociationIndices> reassociation =
             getReassociationForExpansion(indexingMap, expansionInfo);
-        expandedOpOperands.push_back(rewriter.create<TensorExpandShapeOp>(
+        expandedOpOperands.push_back(rewriter.create<tensor::ExpandShapeOp>(
             genericOp.getLoc(), expandedOperandType, opOperand->get(),
             reassociation));
         continue;
       }
     }
     expandedOpOperands.push_back(opOperand->get());
   }
 
   Location loc = genericOp.getLoc();
   SmallVector<Value> outputs;
   for (OpOperand *opOperand : genericOp.getOutputOperands()) {
     AffineMap indexingMap = genericOp.getTiedIndexingMap(opOperand);
     RankedTensorType expandedOutputType =
         getExpandedType(opOperand->get().getType().cast<RankedTensorType>(),
                         indexingMap, expansionInfo);
     if (expandedOutputType != opOperand->get().getType()) {
       SmallVector<ReassociationIndices> reassociation =
           getReassociationForExpansion(indexingMap, expansionInfo);
-      outputs.push_back(rewriter.create<TensorExpandShapeOp>(
+      outputs.push_back(rewriter.create<tensor::ExpandShapeOp>(
           genericOp.getLoc(), expandedOutputType, opOperand->get(),
           reassociation));
     }
   }
 
   // The iterator types of the expanded op are all parallel.
   SmallVector<StringRef> iteratorTypes(expansionInfo.getExpandedOpNumDims(),
                                        getParallelIteratorTypeName());
fuseWithReshapeByExpansion(GenericOp genericOp, Operation *reshapeOp,
   for (OpResult opResult : genericOp->getOpResults()) {
     int64_t resultNumber = opResult.getResultNumber();
     if (!isExpanding && resultTypes[resultNumber] != opResult.getType()) {
       SmallVector<ReassociationIndices> reassociation =
           getReassociationForExpansion(
               genericOp.getTiedIndexingMap(
                   genericOp.getOutputOperand(resultNumber)),
               expansionInfo);
-      resultVals.push_back(rewriter.create<TensorCollapseShapeOp>(
+      resultVals.push_back(rewriter.create<tensor::CollapseShapeOp>(
           genericOp.getLoc(), opResult.getType(),
           fusedOp->getResult(resultNumber), reassociation));
     } else {
       resultVals.push_back(fusedOp->getResult(resultNumber));
     }
   }
   // Assuming a single result.
   return resultVals;
LogicalResult matchAndRewrite(GenericOp genericOp,
     // Only support identity output maps. It could be extended to permuations if
     // needed.
     if (llvm::any_of(genericOp.getOutputOperands(), [&](OpOperand *opOperand) {
           return !genericOp.getTiedIndexingMap(opOperand).isIdentity();
         }))
       return failure();
     int64_t destRank = genericOp.getNumParallelLoops();
     SmallVector<Value> newOperands = genericOp.getInputOperands();
-    TensorExpandShapeOp reshapeFound;
+    tensor::ExpandShapeOp reshapeFound;
     // 1. Look for tensor_expand_shape operands and figure out save the
     // dimensions merged.
     SmallVector<OpOperand *> inputOperands = genericOp.getInputOperands();
     for (auto en : llvm::enumerate(inputOperands)) {
       auto reshapeOp =
-          en.value()->get().template getDefiningOp<TensorExpandShapeOp>();
+          en.value()->get().template getDefiningOp<tensor::ExpandShapeOp>();
       if (!reshapeOp)
         continue;
       // TODO: We could support non-identity map as long as the merged
       // dimensions are still contiguous.
       if (!genericOp.getTiedIndexingMap(en.value()).isIdentity())
         continue;
       if (reshapeFound) {
         // Only support a second reshape op if it has the same reassociate maps.
LogicalResult matchAndRewrite(GenericOp genericOp,
 
     // 4. Reshape the output tensors.
     SmallVector<Value> newOutputs;
     SmallVector<Type> newOutputTypes;
     for (auto output : genericOp.outputs()) {
       auto newOutputType = RankedTensorType::get(
           reshapeFound.getSrcType().getShape(),
           output.getType().template cast<RankedTensorType>().getElementType());
-      Value newOutput = rewriter.create<TensorCollapseShapeOp>(
+      Value newOutput = rewriter.create<tensor::CollapseShapeOp>(
           genericOp->getLoc(), newOutputType, output, reassociation);
       newOutputTypes.push_back(newOutputType);
       newOutputs.push_back(newOutput);
     }
     // 5. Create a new generic op with lowerer rank.
     SmallVector<StringRef> iteratorTypes(newRank,
                                          getParallelIteratorTypeName());
     auto newOp = rewriter.create<GenericOp>(genericOp->getLoc(), newOutputTypes,
                                             newOperands, newOutputs, newMaps,
                                             iteratorTypes);
     rewriter.inlineRegionBefore(genericOp.region(), newOp.region(),
                                 newOp.region().begin());
     // 6. Reshape the so that the type matches the uses.
     SmallVector<Value> newResults;
     for (auto result : llvm::enumerate(newOp->getResults())) {
-      newResults.push_back(rewriter.create<TensorExpandShapeOp>(
+      newResults.push_back(rewriter.create<tensor::ExpandShapeOp>(
           genericOp->getLoc(), genericOp.getOutputTensorTypes()[result.index()],
           result.value(), reassociation));
     }
     rewriter.replaceOp(genericOp, newResults);
     return success();
   }
 };
 
 /// Pattern to fuse a tensor_collapse_shape op with its consumer generic op,
 /// when the reshape op is collapsing dimensions. The dimensionality of the loop
 /// in the consumer is expanded.
 class FoldWithProducerReshapeOpByExpansion
     : public OpRewritePattern<GenericOp> {
 public:
   FoldWithProducerReshapeOpByExpansion(
       MLIRContext *context, ControlElementwiseOpsFusionFn foldReshapes,
       PatternBenefit benefit = 1)
       : OpRewritePattern<GenericOp>(context, benefit),
         controlFoldingReshapes(foldReshapes) {}
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     for (OpOperand *opOperand : genericOp.getInputTensorOperands()) {
-      TensorCollapseShapeOp reshapeOp =
-          opOperand->get().getDefiningOp<TensorCollapseShapeOp>();
+      tensor::CollapseShapeOp reshapeOp =
+          opOperand->get().getDefiningOp<tensor::CollapseShapeOp>();
       if (!reshapeOp)
         continue;
       // Fold only if
       // - The tensor reshape op is folding.
       // - All constraints of fusing with reshape by expansion are met.
       if (!isFusableWithReshapeByDimExpansion(genericOp, opOperand) ||
           (!controlFoldingReshapes(reshapeOp->getResult(0), *opOperand)))
         continue;
LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
     rewriter.replaceOp(reshapeOp, fusedOp->getResults());
     return success();
   }
 };
 
 /// Pattern to fold a tensor_expand_shape op with its producer generic op
 /// by expanding the dimensionality of the loop in the producer op.
 struct FoldReshapeWithGenericOpByExpansion
-    : public OpRewritePattern<TensorExpandShapeOp> {
+    : public OpRewritePattern<tensor::ExpandShapeOp> {
 
   FoldReshapeWithGenericOpByExpansion(
       MLIRContext *context, ControlElementwiseOpsFusionFn foldReshapes,
       PatternBenefit benefit = 1)
-      : OpRewritePattern<TensorExpandShapeOp>(context, benefit),
+      : OpRewritePattern<tensor::ExpandShapeOp>(context, benefit),
         controlFoldingReshapes(foldReshapes) {}
 
-  LogicalResult matchAndRewrite(TensorExpandShapeOp reshapeOp,
+  LogicalResult matchAndRewrite(tensor::ExpandShapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     // Fold only if all constraints of fusing with reshape by expansion are met.
     GenericOp producer = reshapeOp.src().getDefiningOp<GenericOp>();
     if (!producer || producer.getNumOutputs() != 1 ||
         !isFusableWithReshapeByDimExpansion(producer,
                                             producer.getOutputOperand(0)) ||
         !controlFoldingReshapes(producer->getResult(0),
                                 reshapeOp->getOpOperand(0)))
fuseElementwiseOps(PatternRewriter &rewriter, OpOperand *consumerOpOperand,
 
   return fuseElementwiseOpsImpl(producer, consumerOpOperand, controlFn,
                                 rewriter);
 }
 
 bool mlir::linalg::skipUnitDimReshape(const OpResult &producer,
                                       OpOperand &consumer) {
   if (auto producerCollapseOp =
-          dyn_cast<linalg::TensorCollapseShapeOp>(producer.getOwner())) {
+          dyn_cast<tensor::CollapseShapeOp>(producer.getOwner())) {
     return !isUnitDimExpansionOnly(producerCollapseOp);
   }
   if (auto consumerExpandOp =
-          dyn_cast<linalg::TensorExpandShapeOp>(consumer.getOwner())) {
+          dyn_cast<tensor::ExpandShapeOp>(consumer.getOwner())) {
     return !isUnitDimExpansionOnly(consumerExpandOp);
   }
   return true;
 }
 
 namespace {
 /// Patterns to fuse a generic op, with the producer of its operands.
 class FuseElementwiseOps : public OpRewritePattern<GenericOp> {
struct RemoveOutsDependency : public OpRewritePattern<GenericOp> {
   }
 };
 
 } // namespace
 
 void mlir::linalg::populateFoldReshapeOpsByLinearizationPatterns(
     RewritePatternSet &patterns) {
   patterns
-      .add<FoldProducerReshapeOpByLinearization<false, TensorCollapseShapeOp>,
-           FoldProducerReshapeOpByLinearization<false, TensorExpandShapeOp>,
-           FoldConsumerReshapeOpByLinearization<false, TensorCollapseShapeOp>,
-           FoldConsumerReshapeOpByLinearization<false, TensorExpandShapeOp>>(
+      .add<FoldProducerReshapeOpByLinearization<false, tensor::CollapseShapeOp>,
+           FoldProducerReshapeOpByLinearization<false, tensor::ExpandShapeOp>,
+           FoldConsumerReshapeOpByLinearization<false, tensor::CollapseShapeOp>,
+           FoldConsumerReshapeOpByLinearization<false, tensor::ExpandShapeOp>>(
           patterns.getContext());
 }
 
 void mlir::linalg::populateFoldUnitDimsReshapeOpsByLinearizationPatterns(
     RewritePatternSet &patterns) {
   patterns
-      .add<FoldProducerReshapeOpByLinearization<true, TensorCollapseShapeOp>,
-           FoldProducerReshapeOpByLinearization<true, TensorExpandShapeOp>,
-           FoldConsumerReshapeOpByLinearization<true, TensorCollapseShapeOp>,
-           FoldConsumerReshapeOpByLinearization<true, TensorExpandShapeOp>>(
+      .add<FoldProducerReshapeOpByLinearization<true, tensor::CollapseShapeOp>,
+           FoldProducerReshapeOpByLinearization<true, tensor::ExpandShapeOp>,
+           FoldConsumerReshapeOpByLinearization<true, tensor::CollapseShapeOp>,
+           FoldConsumerReshapeOpByLinearization<true, tensor::ExpandShapeOp>>(
           patterns.getContext());
 }
 
 void mlir::linalg::populateFoldReshapeOpsByExpansionPatterns(
     RewritePatternSet &patterns,
     ControlElementwiseOpsFusionFn controlFoldingReshapes) {
   patterns.add<FoldReshapeWithGenericOpByExpansion>(patterns.getContext(),
                                                     controlFoldingReshapes);
   patterns.add<FoldWithProducerReshapeOpByExpansion>(patterns.getContext(),
                                                      controlFoldingReshapes);
 }
 
 void mlir::linalg::populateElementwiseOpsFusionPatterns(
     RewritePatternSet &patterns, LinalgElementwiseFusionOptions options) {
   auto *context = patterns.getContext();
   patterns.add<FuseElementwiseOps, FoldScalarOrSplatConstant,
                FoldConstantTranspose>(context,
                                       options.controlElementwiseOpsFusionFn);
   patterns.add<RemoveOutsDependency>(context);
   populateFoldReshapeOpsByExpansionPatterns(patterns,
                                             options.controlFoldingReshapesFn);
   AffineApplyOp::getCanonicalizationPatterns(patterns, context);
   GenericOp::getCanonicalizationPatterns(patterns, context);
-  TensorExpandShapeOp::getCanonicalizationPatterns(patterns, context);
-  TensorCollapseShapeOp::getCanonicalizationPatterns(patterns, context);
+  tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
+  tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
   context->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(
       patterns);
 }
 
 void mlir::linalg::populatePushReshapeOpsPatterns(RewritePatternSet &patterns) {
   auto *context = patterns.getContext();
   patterns.add<PushExpandingReshape>(context);
 }
 
 std::unique_ptr<Pass> mlir::createLinalgElementwiseOpFusionPass() {
   return std::make_unique<LinalgElementwiseOpFusionPass>();
 }
 
 std::unique_ptr<Pass> mlir::createFoldReshapeOpsByLinearizationPass() {
   return std::make_unique<FoldReshapeOpsByLinearizationPass>();
 }

mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp

 //===- ResolveShapedTypeResultDims.cpp - Resolve dim ops of result values -===//
 //
 // 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 pass resolves `memref.dim` operations of result values in terms of
 // shapes of their operands using the `InferShapedTypeOpInterface`.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 

mlir/lib/Dialect/Tensor/IR/CMakeLists.txt

+set(LLVM_OPTIONAL_SOURCES
+  TensorDialect.cpp
+  TensorInferTypeOpInterfaceImpl.cpp
+  TensorOps.cpp
+)
+
 add_mlir_dialect_library(MLIRTensor
   TensorDialect.cpp
   TensorOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/include/mlir/Dialect/Tensor
 
   DEPENDS
   MLIRTensorOpsIncGen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRArithmetic
   MLIRCastInterfaces
   MLIRDialectUtils
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRSideEffectInterfaces
   MLIRSupport
   MLIRStandard
   MLIRViewLikeInterface
   )
+
+add_mlir_dialect_library(MLIRTensorInferTypeOpInterfaceImpl
+  TensorInferTypeOpInterfaceImpl.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor
+
+  LINK_LIBS PUBLIC
+  MLIRAffine
+  MLIRIR
+  MLIRInferTypeOpInterface
+  MLIRStandard
+  MLIRSupport
+  MLIRTensor
+  )

mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp

+//===- InferTypeOpImpl.cpp - InferType Interface external models *- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/StandardOps/Utils/Utils.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+
+using namespace mlir;
+using namespace mlir::tensor;
+
+/// Compute a map that for a given dimension of the expanded type gives the
+/// dimension in the collapsed type it maps to. Essentially its the inverse of
+/// the `reassocation` maps.
+static llvm::DenseMap<int64_t, int64_t>
+getExpandedDimToCollapsedDimMap(ArrayRef<AffineMap> reassociation) {
+  llvm::DenseMap<int64_t, int64_t> expandedDimToCollapsedDim;
+  for (const auto &map : enumerate(reassociation)) {
+    unsigned startPos =
+        map.value().getResults().front().cast<AffineDimExpr>().getPosition();
+    unsigned endPos =
+        map.value().getResults().back().cast<AffineDimExpr>().getPosition();
+    for (auto dim : llvm::seq_inclusive(startPos, endPos)) {
+      expandedDimToCollapsedDim[dim] = map.index();
+    }
+  }
+  return expandedDimToCollapsedDim;
+}
+
+/// For reshape op compute the shape at dimension `dimIndex` of the output in
+/// terms of shape of the `src`, when the reshape op is a collapsing
+/// operation. It is the product of the shape of the collapsed dimensions of the
+/// `src`.
+static OpFoldResult
+getCollapsedOutputDimFromInputShape(OpBuilder &builder, Location loc,
+                                    int64_t dimIndex, Value src,
+                                    ArrayRef<AffineMap> reassociationMap) {
+  AffineMap map = reassociationMap[dimIndex];
+  unsigned startPos =
+      map.getResults().front().cast<AffineDimExpr>().getPosition();
+  unsigned endPos = map.getResults().back().cast<AffineDimExpr>().getPosition();
+  AffineExpr expr;
+  SmallVector<Value, 2> dynamicDims;
+  for (auto dim : llvm::seq_inclusive(startPos, endPos)) {
+    dynamicDims.push_back(builder.createOrFold<tensor::DimOp>(loc, src, dim));
+    AffineExpr currExpr = builder.getAffineSymbolExpr(dim - startPos);
+    expr = (expr ? expr * currExpr : currExpr);
+  }
+  return applyMapToValues(builder, loc,
+                          AffineMap::get(0, endPos - startPos + 1, expr),
+                          dynamicDims)[0];
+}
+
+/// Given the `src` of a collapsing reshape op and its reassociation maps,
+/// compute the shape of the result of the reshape.
+static SmallVector<OpFoldResult, 4> getCollapsedOutputShapeFromInputShape(
+    OpBuilder &builder, Location loc, Value src,
+    ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
+  return llvm::to_vector<4>(llvm::map_range(
+      llvm::seq<int64_t>(0, dstStaticShape.size()), [&](int64_t dim) {
+        return getCollapsedOutputDimFromInputShape(builder, loc, dim, src,
+                                                   reassociation);
+      }));
+}
+
+/// For an expanding reshape op, compute the value for a dimension of the output
+/// from the shape of the input.
+static OpFoldResult getExpandedOutputDimFromInputShape(
+    OpBuilder &builder, Location loc, int64_t dimIndex, Value src,
+    ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation,
+    llvm::DenseMap<int64_t, int64_t> &expandedDimToCollapsedDim) {
+  if (!ShapedType::isDynamic(dstStaticShape[dimIndex])) {
+    return builder.getI64IntegerAttr(dstStaticShape[dimIndex]);
+  }
+  unsigned sourceDimPos = expandedDimToCollapsedDim[dimIndex];
+  unsigned startPos = reassociation[sourceDimPos]
+                          .getResults()
+                          .front()
+                          .cast<AffineDimExpr>()
+                          .getPosition();
+  unsigned endPos = reassociation[sourceDimPos]
+                        .getResults()
+                        .back()
+                        .cast<AffineDimExpr>()
+                        .getPosition();
+  int64_t linearizedStaticDim = 1;
+  for (auto &d :
+       llvm::enumerate(dstStaticShape.slice(startPos, endPos - startPos + 1))) {
+    if (d.index() + startPos == static_cast<unsigned>(dimIndex))
+      continue;
+    assert(!ShapedType::isDynamic(d.value()) &&
+           "single dimension cannot be expanded into multiple dynamic "
+           "dimensions");
+    linearizedStaticDim *= d.value();
+  }
+  Value sourceDim = builder.create<tensor::DimOp>(loc, src, sourceDimPos);
+  return applyMapToValues(
+      builder, loc,
+      AffineMap::get(
+          0, 1, builder.getAffineSymbolExpr(0).floorDiv(linearizedStaticDim)),
+      sourceDim)[0];
+}
+
+/// Given the `src` of an expanding reshape op, the reassociation maps and the
+/// result type, compute the shape of the result of the reshape.
+static SmallVector<OpFoldResult, 4> getExpandedOutputShapeFromInputShape(
+    OpBuilder &builder, Location loc, Value src,
+    ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
+  llvm::DenseMap<int64_t, int64_t> expandedDimToCollapsedDim =
+      getExpandedDimToCollapsedDimMap(reassociation);
+  return llvm::to_vector<4>(llvm::map_range(
+      llvm::seq<int64_t>(0, dstStaticShape.size()), [&](int64_t dim) {
+        return getExpandedOutputDimFromInputShape(builder, loc, dim, src,
+                                                  dstStaticShape, reassociation,
+                                                  expandedDimToCollapsedDim);
+      }));
+}
+
+static SmallVector<OpFoldResult, 4>
+getReshapeOutputShapeFromInputShape(OpBuilder &builder, Location loc, Value src,
+                                    ArrayRef<int64_t> dstStaticShape,
+                                    ArrayRef<AffineMap> reassocation) {
+  return dstStaticShape.size() >
+                 static_cast<size_t>(src.getType().cast<ShapedType>().getRank())
+             ? getExpandedOutputShapeFromInputShape(
+                   builder, loc, src, dstStaticShape, reassocation)
+             : getCollapsedOutputShapeFromInputShape(
+                   builder, loc, src, dstStaticShape, reassocation);
+}
+
+/// Helper function to convert a vector of `OpFoldResult`s into a vector of
+/// `Value`s.
+static SmallVector<Value> getAsValues(OpBuilder &b, Location loc,
+                                      ArrayRef<OpFoldResult> valueOrAttrVec) {
+  return llvm::to_vector<4>(
+      llvm::map_range(valueOrAttrVec, [&](OpFoldResult value) -> Value {
+        return getValueOrCreateConstantIndexOp(b, loc, value);
+      }));
+}
+
+template <typename OpTy>
+struct ReifyExpandOrCollapseShapeOp
+    : public ReifyRankedShapedTypeOpInterface::ExternalModel<
+          ReifyExpandOrCollapseShapeOp<OpTy>, OpTy> {
+  LogicalResult
+  reifyResultShapes(Operation *op, OpBuilder &b,
+                    ReifiedRankedShapedTypeDims &reifiedReturnShapes) const {
+    auto loc = op->getLoc();
+    auto reshape_op = cast<OpTy>(op);
+    auto result_shape = getReshapeOutputShapeFromInputShape(
+        b, loc, reshape_op.src(), reshape_op.getResultType().getShape(),
+        reshape_op.getReassociationMaps());
+    reifiedReturnShapes.push_back(getAsValues(b, loc, result_shape));
+    return success();
+  }
+};
+
+void mlir::tensor::registerInferTypeOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry
+      .addOpInterface<tensor::ExpandShapeOp,
+                      ReifyExpandOrCollapseShapeOp<tensor::ExpandShapeOp>>();
+  registry
+      .addOpInterface<tensor::CollapseShapeOp,
+                      ReifyExpandOrCollapseShapeOp<tensor::CollapseShapeOp>>();
+}

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

 //===----------------------------------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/StandardOps/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/STLExtras.h"
if (resultRankedType) {
     if (shapeSize != resultRankedType.getRank())
       return op.emitOpError(
           "length of shape operand differs from the result's tensor rank");
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// Reassociative reshape ops
+//===----------------------------------------------------------------------===//
+
+SmallVector<AffineMap, 4> CollapseShapeOp::getReassociationMaps() {
+  return getSymbolLessAffineMaps(getReassociationExprs());
+}
+SmallVector<ReassociationExprs, 4> CollapseShapeOp::getReassociationExprs() {
+  return convertReassociationIndicesToExprs(getContext(),
+                                            getReassociationIndices());
+}
+
+SmallVector<AffineMap, 4> ExpandShapeOp::getReassociationMaps() {
+  return getSymbolLessAffineMaps(getReassociationExprs());
+}
+SmallVector<ReassociationExprs, 4> ExpandShapeOp::getReassociationExprs() {
+  return convertReassociationIndicesToExprs(getContext(),
+                                            getReassociationIndices());
+}
+
+static void print(OpAsmPrinter &p, ExpandShapeOp op) {
+  ::mlir::printReshapeOp<ExpandShapeOp>(p, op);
+}
+
+static void print(OpAsmPrinter &p, CollapseShapeOp op) {
+  ::mlir::printReshapeOp<CollapseShapeOp>(p, op);
+}
+
+/// Compute the RankedTensorType obtained by applying `reassociation` to `type`.
+static RankedTensorType
+computeTensorReshapeCollapsedType(RankedTensorType type,
+                                  ArrayRef<AffineMap> reassociation) {
+  auto shape = type.getShape();
+  SmallVector<int64_t, 4> newShape;
+  newShape.reserve(reassociation.size());
+
+  // Use the fact that reassociation is valid to simplify the logic: only use
+  // each map's rank.
+  assert(isReassociationValid(reassociation) && "invalid reassociation");
+  unsigned currentDim = 0;
+  for (AffineMap m : reassociation) {
+    unsigned dim = m.getNumResults();
+    auto band = shape.slice(currentDim, dim);
+    int64_t size = 1;
+    if (llvm::is_contained(band, ShapedType::kDynamicSize))
+      size = ShapedType::kDynamicSize;
+    else
+      for (unsigned d = 0; d < dim; ++d)
+        size *= shape[currentDim + d];
+    newShape.push_back(size);
+    currentDim += dim;
+  }
+
+  return RankedTensorType::get(newShape, type.getElementType());
+}
+
+void CollapseShapeOp::build(OpBuilder &b, OperationState &result, Value src,
+                            ArrayRef<ReassociationIndices> reassociation,
+                            ArrayRef<NamedAttribute> attrs) {
+  auto resultType = computeTensorReshapeCollapsedType(
+      src.getType().cast<RankedTensorType>(),
+      getSymbolLessAffineMaps(
+          convertReassociationIndicesToExprs(b.getContext(), reassociation)));
+  build(b, result, resultType, src, attrs);
+  result.addAttribute(getReassociationAttrName(),
+                      getReassociationIndicesAttribute(b, reassociation));
+}
+
+void ExpandShapeOp::build(OpBuilder &b, OperationState &result, Value src,
+                          ArrayRef<ReassociationIndices> reassociation,
+                          ArrayRef<NamedAttribute> attrs) {
+  auto resultType = computeTensorReshapeCollapsedType(
+      src.getType().cast<RankedTensorType>(),
+      getSymbolLessAffineMaps(
+          convertReassociationIndicesToExprs(b.getContext(), reassociation)));
+  build(b, result, resultType, src, attrs);
+  result.addAttribute(getReassociationAttrName(),
+                      getReassociationIndicesAttribute(b, reassociation));
+}
+
+template <typename TensorReshapeOp, bool isExpansion = std::is_same<
+                                        TensorReshapeOp, ExpandShapeOp>::value>
+static LogicalResult verifyTensorReshapeOp(TensorReshapeOp op,
+                                           RankedTensorType expandedType,
+                                           RankedTensorType collapsedType) {
+  if (failed(
+          verifyReshapeLikeTypes(op, expandedType, collapsedType, isExpansion)))
+    return failure();
+
+  auto maps = op.getReassociationMaps();
+  RankedTensorType expectedType =
+      computeTensorReshapeCollapsedType(expandedType, maps);
+  if (collapsedType != expectedType)
+    return op.emitOpError("expected collapsed type to be ")
+           << expectedType << ", but got " << collapsedType;
+  return success();
+}
+
+static LogicalResult verify(ExpandShapeOp op) {
+  return verifyTensorReshapeOp(op, op.getResultType(), op.getSrcType());
+}
+
+static LogicalResult verify(CollapseShapeOp op) {
+  return verifyTensorReshapeOp(op, op.getSrcType(), op.getResultType());
+}
+
+namespace {
+/// Reshape of a splat constant can be replaced with a constant of the result
+/// type.
+template <typename TensorReshapeOp>
+struct FoldReshapeWithConstant : OpRewritePattern<TensorReshapeOp> {
+  using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
+                                PatternRewriter &rewriter) const override {
+    DenseElementsAttr attr;
+    if (!matchPattern(reshapeOp.src(), m_Constant(&attr)))
+      return failure();
+    if (!attr || !attr.isSplat())
+      return failure();
+    DenseElementsAttr newAttr = DenseElementsAttr::getFromRawBuffer(
+        reshapeOp.getResultType(), attr.getRawData(), true);
+    rewriter.replaceOpWithNewOp<arith::ConstantOp>(reshapeOp, newAttr);
+    return success();
+  }
+};
+
+} // namespace
+
+void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                                MLIRContext *context) {
+  results.add<CollapseReshapeOps<ExpandShapeOp>,
+              CollapseMixedReshapeOps<ExpandShapeOp, CollapseShapeOp>,
+              FoldReshapeWithConstant<ExpandShapeOp>>(context);
+}
+
+void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                                  MLIRContext *context) {
+  results.add<CollapseReshapeOps<CollapseShapeOp>,
+              CollapseMixedReshapeOps<CollapseShapeOp, ExpandShapeOp>,
+              FoldReshapeWithConstant<CollapseShapeOp>>(context);
+}
+
+OpFoldResult ExpandShapeOp::fold(ArrayRef<Attribute> operands) {
+  return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this, operands);
+}
+OpFoldResult CollapseShapeOp::fold(ArrayRef<Attribute> operands) {
+  return foldReshapeOp<CollapseShapeOp, ExpandShapeOp>(*this, operands);
+}
+
+//===----------------------------------------------------------------------===//
 // ExtractSliceOp
 //===----------------------------------------------------------------------===//
 
 /// An extract_slice op result type can be fully inferred from the source type
 /// and the static representation of offsets, sizes and strides. Special
 /// sentinels encode the dynamic case.
 RankedTensorType
 ExtractSliceOp::inferResultType(RankedTensorType sourceRankedTensorType,

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

 
 
 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> ()>
 // CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL: @test_broadcast
 func @test_broadcast(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2] : tensor<2xf32>
-  // CHECK: [[RESHAPE:%.+]] = linalg.tensor_collapse_shape %arg0
+  // CHECK: [[RESHAPE:%.+]] = tensor.collapse_shape %arg0
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel"]} ins([[RESHAPE]], %arg1 : tensor<f32>, tensor<2xf32>) outs([[INIT]] : tensor<2xf32>) {
   // CHECK: ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
   // CHECK:   [[ELEMENT:%.+]] = arith.addf %arg2, %arg3 : f32
   // CHECK:   linalg.yield [[ELEMENT]] : f32
   // CHECK: } -> tensor<2xf32>
   %0 = "tosa.add"(%arg0, %arg1) : (tensor<1xf32>, tensor<2xf32>) -> tensor<2xf32>
   return %0 : tensor<2xf32>
 }
 
 // -----
 
 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
 // CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> ()>
 
 // CHECK-LABEL: @test_broadcast_swapped_args
 func @test_broadcast_swapped_args(%arg0: tensor<2xf32>, %arg1: tensor<1xf32>) -> tensor<2xf32> {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2] : tensor<2xf32>
-  // CHECK: [[RESHAPE:%.+]] = linalg.tensor_collapse_shape %arg1
+  // CHECK: [[RESHAPE:%.+]] = tensor.collapse_shape %arg1
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0, [[RESHAPE]] : tensor<2xf32>, tensor<f32>) outs([[INIT]] : tensor<2xf32>) {
   // CHECK: ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
   // CHECK:   [[ELEMENT:%.+]] = arith.addf %arg2, %arg3 : f32
   // CHECK:   linalg.yield [[ELEMENT]] : f32
   // CHECK: } -> tensor<2xf32>
   %0 = "tosa.add"(%arg0, %arg1) : (tensor<2xf32>, tensor<1xf32>) -> tensor<2xf32>
   return %0 : tensor<2xf32>
 }
 
 // -----
 
 // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
 // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
 
 // CHECK-LABEL: @test_multibroadcast
 func @test_multibroadcast(%arg0: tensor<1x3xf32>, %arg1: tensor<2x1xf32>) -> tensor<2x3xf32> {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 3] : tensor<2x3xf32>
-  // CHECK: [[RESHAPE1:%.+]] = linalg.tensor_collapse_shape %arg0 {{\[}}[0, 1]]
-  // CHECK: [[RESHAPE2:%.+]] = linalg.tensor_collapse_shape %arg1 {{\[}}[0, 1]]
+  // CHECK: [[RESHAPE1:%.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1]]
+  // CHECK: [[RESHAPE2:%.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1]]
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP0]]], iterator_types = ["parallel", "parallel"]} ins([[RESHAPE1]], [[RESHAPE2]] : tensor<3xf32>, tensor<2xf32>) outs([[INIT]] : tensor<2x3xf32>) {
   // CHECK: ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
   // CHECK:   [[ELEMENT:%.+]] = arith.addf %arg2, %arg3 : f32
   // CHECK:   linalg.yield [[ELEMENT]] : f32
   // CHECK: } -> tensor<2x3xf32>
   %0 = "tosa.add"(%arg0, %arg1) : (tensor<1x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32>
   return %0 : tensor<2x3xf32>
 }
func @test_negate_quantized(%arg0: tensor<1xi8>) -> () {
 
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_reshape_downrank
 func @test_reshape_downrank(%arg0: tensor<2x3xf32>) -> tensor<6xf32> {
-  // CHECK: [[RESHAPE:%.+]] = linalg.tensor_collapse_shape %arg0 {{\[}}[0, 1]]
+  // CHECK: [[RESHAPE:%.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1]]
   %0 = "tosa.reshape"(%arg0) {new_shape = [6]} : (tensor<2x3xf32>) -> tensor<6xf32>
   // CHECK: return [[RESHAPE]]
   return %0 : tensor<6xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @test_reshape_downrank_dyn
 func @test_reshape_downrank_dyn(%arg0: tensor<2x?xf32>) -> tensor<?xf32> {
-  // CHECK: [[RESHAPE:%.+]] = linalg.tensor_collapse_shape %arg0 {{\[}}[0, 1]]
+  // CHECK: [[RESHAPE:%.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1]]
   %0 = "tosa.reshape"(%arg0) {new_shape = [-1]} : (tensor<2x?xf32>) -> tensor<?xf32>
   // CHECK: return [[RESHAPE]]
   return %0 : tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @test_reshape_uprank
 func @test_reshape_uprank(%arg0: tensor<6xf32>) -> tensor<2x3xf32> {
-  // CHECK: [[RESHAPE:%.+]] = linalg.tensor_expand_shape %arg0 {{\[}}[0, 1]]
+  // CHECK: [[RESHAPE:%.+]] = tensor.expand_shape %arg0 {{\[}}[0, 1]]
   %0 = "tosa.reshape"(%arg0) {new_shape = [2, 3]} : (tensor<6xf32>) -> tensor<2x3xf32>
   // CHECK: return [[RESHAPE]]
   return %0 : tensor<2x3xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @test_reshape_uprank_dyn
 func @test_reshape_uprank_dyn(%arg0: tensor<?xf32>) -> tensor<2x?xf32> {
-  // CHECK: [[RESHAPE:%.+]] = linalg.tensor_expand_shape %arg0 {{\[}}[0, 1]]
+  // CHECK: [[RESHAPE:%.+]] = tensor.expand_shape %arg0 {{\[}}[0, 1]]
   %0 = "tosa.reshape"(%arg0) {new_shape = [2, -1]} : (tensor<?xf32>) -> tensor<2x?xf32>
   // CHECK: return [[RESHAPE]]
   return %0 : tensor<2x?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @test_reshape_samerank
 func @test_reshape_samerank(%arg0: tensor<3x2xf32>) -> tensor<2x3xf32> {
   // CHECK-SAME: (%[[ARG0:.*]]: tensor<3x2xf32>)
-  // CHECK-NEXT: %[[RESHAPE1:.*]] = linalg.tensor_collapse_shape %[[ARG0]] {{\[}}[0, 1]]
-  // CHECK-NEXT: %[[RESHAPE2:.*]] = linalg.tensor_expand_shape %[[RESHAPE1]] {{\[}}[0, 1]]
+  // CHECK-NEXT: %[[RESHAPE1:.*]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1]]
+  // CHECK-NEXT: %[[RESHAPE2:.*]] = tensor.expand_shape %[[RESHAPE1]] {{\[}}[0, 1]]
   %0 = "tosa.reshape"(%arg0) {new_shape = [2, 3]} : (tensor<3x2xf32>) -> tensor<2x3xf32>
   // CHECK-NEXT: return %[[RESHAPE2]]
   return %0 : tensor<2x3xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @test_reshape_samerank_dyn
 func @test_reshape_samerank_dyn(%arg0: tensor<?x2xf32>) -> tensor<2x?xf32> {
   // CHECK-SAME: (%[[ARG0:.*]]: tensor<?x2xf32>)
-  // CHECK-NEXT: %[[RESHAPE1:.*]] = linalg.tensor_collapse_shape %[[ARG0]] {{\[}}[0, 1]]
-  // CHECK-NEXT: %[[RESHAPE2:.*]] = linalg.tensor_expand_shape %[[RESHAPE1]] {{\[}}[0, 1]]
+  // CHECK-NEXT: %[[RESHAPE1:.*]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1]]
+  // CHECK-NEXT: %[[RESHAPE2:.*]] = tensor.expand_shape %[[RESHAPE1]] {{\[}}[0, 1]]
   %0 = "tosa.reshape"(%arg0) {new_shape = [2, -1]} : (tensor<?x2xf32>) -> tensor<2x?xf32>
   // CHECK-NEXT: return %[[RESHAPE2]]
   return %0 : tensor<2x?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @test_reshape_downrank_6D
 func @test_reshape_downrank_6D(%arg0: tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32> {
-  // CHECK: linalg.tensor_collapse_shape %arg0 {{\[}}[0, 1, 2], [3], [4, 5]]
+  // CHECK: tensor.collapse_shape %arg0 {{\[}}[0, 1, 2], [3], [4, 5]]
   %0 = "tosa.reshape"(%arg0) {new_shape = [6, 5, 77]} : (tensor<1x2x3x5x7x11xf32>) -> tensor<6x5x77xf32>
   return %0 : tensor<6x5x77xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @test_reshape_downrank_6D_dyn
 func @test_reshape_downrank_6D_dyn(%arg0: tensor<1x2x?x5x7x11xf32>) -> tensor<?x5x77xf32> {
-  // CHECK: linalg.tensor_collapse_shape %arg0 {{\[}}[0, 1, 2, 3, 4, 5]]
-  // CHECK: linalg.tensor_expand_shape %0 {{\[}}[0, 1, 2]]
+  // CHECK: tensor.collapse_shape %arg0 {{\[}}[0, 1, 2, 3, 4, 5]]
+  // CHECK: tensor.expand_shape %0 {{\[}}[0, 1, 2]]
   %0 = "tosa.reshape"(%arg0) {new_shape = [-1, 5, 77]} : (tensor<1x2x?x5x7x11xf32>) -> tensor<?x5x77xf32>
   return %0 : tensor<?x5x77xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @test_identity
 func @test_identity(%arg0: tensor<1xf32>, %arg1: tensor<1xi32>) -> (tensor<1xf32>, tensor<1xi32>) {

 func @reduce_float(%arg0: tensor<5x4xf32>) -> () {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [4]
   // CHECK: [[CST0:%.+]] = arith.constant 0.0
   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<4xf32>)
   // CHECK: ^bb0(%arg1: f32, %arg2: f32)
   // CHECK:   [[RES:%.+]] = arith.addf %arg1, %arg2 : f32
   // CHECK:   linalg.yield [[RES]] : f32
-  // CHECK: linalg.tensor_expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xf32> into tensor<1x4xf32>
+  // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xf32> into tensor<1x4xf32>
   %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<1x4xf32>
 
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [5]
   // CHECK: [[CST0:%.+]] = arith.constant 0.0
   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<5xf32>)
   // CHECK: ^bb0(%arg1: f32, %arg2: f32)
   // CHECK:   [[RES:%.+]] = arith.addf %arg1, %arg2 : f32
   // CHECK:   linalg.yield [[RES]] : f32
-  // CHECK: linalg.tensor_expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<5xf32> into tensor<5x1xf32>
+  // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<5xf32> into tensor<5x1xf32>
   %1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xf32>) -> tensor<5x1xf32>
 
   // CHECK: arith.constant 1.0
   // CHECK: linalg.fill
   // CHECK: linalg.generic
   // CHECK: arith.mulf
   %2 = "tosa.reduce_prod"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<1x4xf32>
 

 func @reduce_int(%arg0: tensor<5x4xi32>) -> () {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [4]
   // CHECK: [[CST0:%.+]] = arith.constant 0
   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<4xi32>)
   // CHECK: ^bb0(%arg1: i32, %arg2: i32)
   // CHECK:   [[RES:%.+]] = arith.addi %arg1, %arg2 : i32
   // CHECK:   linalg.yield [[RES]] : i32
-  // CHECK: linalg.tensor_expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xi32> into tensor<1x4xi32>
+  // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xi32> into tensor<1x4xi32>
   %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<1x4xi32>
 
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [5]
   // CHECK: [[CST0:%.+]] = arith.constant 0
   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<5xi32>)
   // CHECK: ^bb0(%arg1: i32, %arg2: i32)
   // CHECK:   [[RES:%.+]] = arith.addi %arg1, %arg2 : i32
   // CHECK:   linalg.yield [[RES]] : i32
-  // CHECK: linalg.tensor_expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<5xi32> into tensor<5x1xi32>
+  // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<5xi32> into tensor<5x1xi32>
   %1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xi32>) -> tensor<5x1xi32>
 
   // CHECK: arith.constant 1
   // CHECK: linalg.fill
   // CHECK: linalg.generic
   // CHECK: arith.muli
   %2 = "tosa.reduce_prod"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<1x4xi32>
 

 func @reduce_bool(%arg0: tensor<5x4xi1>) -> () {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [4]
   // CHECK: [[CST0:%.+]] = arith.constant true
   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xi1>) outs([[FILL]] : tensor<4xi1>)
   // CHECK: ^bb0(%arg1: i1, %arg2: i1)
   // CHECK:   [[RES:%.+]] = arith.andi %arg1, %arg2 : i1
   // CHECK:   linalg.yield [[RES]] : i1
-  // CHECK: linalg.tensor_expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xi1> into tensor<1x4xi1>
+  // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xi1> into tensor<1x4xi1>
   %0 = "tosa.reduce_all"(%arg0) {axis = 0 : i64} : (tensor<5x4xi1>) -> tensor<1x4xi1>
 
   // CHECK: arith.constant false
   // CHECK: linalg.fill
   // CHECK: linalg.generic
   // CHECK: or
   %1 = "tosa.reduce_any"(%arg0) {axis = 0 : i64} : (tensor<5x4xi1>) -> tensor<1x4xi1>
 

 // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)>
 // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 
 // CHECK-LABEL: @tile
 func @tile(%arg0 : tensor<2x3xi8>) -> () {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 2, 1, 3]
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs([[INIT]] : tensor<2x2x1x3xi8>)
   // CHECK:   linalg.yield %arg1 : i8
-  // CHECK: linalg.tensor_collapse_shape [[GENERIC]] {{\[}}[0, 1, 2], [3]]
+  // CHECK: tensor.collapse_shape [[GENERIC]] {{\[}}[0, 1, 2], [3]]
   %0 = "tosa.tile"(%arg0) {multiples = [2, 1]} : (tensor<2x3xi8>)  -> (tensor<4x3xi8>)
 
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 2, 2, 3]
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs([[INIT]] : tensor<1x2x2x3xi8>)
   // CHECK:   linalg.yield %arg1 : i8
-  // CHECK: linalg.tensor_collapse_shape [[GENERIC]] {{\[}}[0, 1], [2, 3]]
+  // CHECK: tensor.collapse_shape [[GENERIC]] {{\[}}[0, 1], [2, 3]]
   %1 = "tosa.tile"(%arg0) {multiples = [1, 2]} : (tensor<2x3xi8>)  -> (tensor<2x6xi8>)
 
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [5, 2, 7, 3]
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs([[INIT]] : tensor<5x2x7x3xi8>)
   // CHECK:   linalg.yield %arg1 : i8
-  // CHECK: linalg.tensor_collapse_shape [[GENERIC]] {{\[}}[0, 1], [2, 3]]
+  // CHECK: tensor.collapse_shape [[GENERIC]] {{\[}}[0, 1], [2, 3]]
   %2 = "tosa.tile"(%arg0) {multiples = [5, 7]} : (tensor<2x3xi8>)  -> (tensor<10x21xi8>)
 
   return
 }
 
 // -----
 
 

 
 // CHECK-LABEL: @depthwise_conv
 func @depthwise_conv(%arg0 : tensor<1x7x5x3xf32>, %arg1 : tensor<3x1x3x11xf32>, %arg2 : tensor<33xf32>) -> () {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 5, 5, 3, 11]
   // CHECK: [[CST0:%.+]] = arith.constant 0
   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
   // CHECK: [[OUT:%.+]] = linalg.init_tensor [1, 5, 5, 33]
   // CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<1x7x5x3xf32>, tensor<3x1x3x11xf32>) outs([[FILL]] : tensor<1x5x5x3x11xf32>)
-  // CHECK: [[COLLAPSED:%.+]] = linalg.tensor_collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
+  // CHECK: [[COLLAPSED:%.+]] = tensor.collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
   // CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<33xf32>, tensor<1x5x5x33xf32>) outs([[OUT]] : tensor<1x5x5x33xf32>) {
   // CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):  // no predecessors
   // CHECK:   [[ADD:%.+]] = arith.addf %arg3, %arg4 : f32
   // CHECK:   linalg.yield [[ADD]] : f32
   // CHECK: } -> tensor<1x5x5x33xf32>
   %2 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) { pad = [0, 0, 0, 0], stride = [1, 1], dilation = [1, 1] } : (tensor<1x7x5x3xf32>, tensor<3x1x3x11xf32>, tensor<33xf32>)  -> (tensor<1x5x5x33xf32>)
   return
 }
 
 // -----
 
 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)>
 // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 
 // CHECK-LABEL: @depthwise_conv_strides
 func @depthwise_conv_strides(%arg0 : tensor<1x11x9x3xf32>, %arg1 : tensor<3x1x3x11xf32>, %arg2 : tensor<33xf32>) -> () {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 5, 5, 3, 11]
   // CHECK: [[CST0:%.+]] = arith.constant 0
   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
   // CHECK: [[OUT:%.+]] = linalg.init_tensor [1, 5, 5, 33]
   // CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<1x11x9x3xf32>, tensor<3x1x3x11xf32>) outs([[FILL]] : tensor<1x5x5x3x11xf32>)
-  // CHECK: [[COLLAPSED:%.+]] = linalg.tensor_collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
+  // CHECK: [[COLLAPSED:%.+]] = tensor.collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
   // CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<33xf32>, tensor<1x5x5x33xf32>) outs([[OUT]] : tensor<1x5x5x33xf32>) {
   // CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):  // no predecessors
   // CHECK:   [[ADD:%.+]] = arith.addf %arg3, %arg4 : f32
   // CHECK:   linalg.yield [[ADD]] : f32
   // CHECK: } -> tensor<1x5x5x33xf32>
   %2 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) { pad = [0, 0, 0, 0], stride = [2, 2], dilation = [1, 1] } : (tensor<1x11x9x3xf32>, tensor<3x1x3x11xf32>, tensor<33xf32>)  -> (tensor<1x5x5x33xf32>)
   return
 }
func @depthwise_conv_quant(%arg0 : tensor<1x12x12x4xi8>, %arg1 : tensor<3x3x4x128xi8>, %arg2 : tensor<512xi32>) -> () {
 
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 12, 12, 4, 128]
   // CHECK: [[CST0:%.+]] = arith.constant 0
   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
   // CHECK: [[OUT:%.+]] = linalg.init_tensor [1, 12, 12, 512]
   // CHECK: [[C128:%.+]] = arith.constant -128
   // CHECK: [[C42:%.+]] = arith.constant 42
   // CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm_q {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins([[PAD]], %arg1, [[C128]], [[C42]] : tensor<1x14x14x4xi8>, tensor<3x3x4x128xi8>, i32, i32) outs([[FILL]] : tensor<1x12x12x4x128xi32>)
-  // CHECK: [[COLLAPSED:%.+]] = linalg.tensor_collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
+  // CHECK: [[COLLAPSED:%.+]] = tensor.collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
   // CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<512xi32>, tensor<1x12x12x512xi32>) outs([[OUT]] : tensor<1x12x12x512xi32>) {
   // CHECK: ^bb0(%arg3: i32, %arg4: i32, %arg5: i32):  // no predecessors
   // CHECK:   [[ADD:%.+]] = arith.addi %arg3, %arg4 : i32
   // CHECK:   linalg.yield [[ADD]] : i32
   // CHECK: } -> tensor<1x12x12x512xi32>
   %0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [1, 1, 1, 1], quantization_info = {input_zp = -128 : i32, weight_zp = 42 : i32}, stride = [1, 1], dilation = [1, 1] } : (tensor<1x12x12x4xi8>, tensor<3x3x4x128xi8>, tensor<512xi32>)  -> tensor<1x12x12x512xi32>
   return
 }
 
 // -----
 
 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)>
 // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 
 // CHECK-LABEL: @depthwise_conv_quant_dilations
 func @depthwise_conv_quant_dilations(%arg0 : tensor<1x14x14x4xi8>, %arg1 : tensor<3x3x4x128xi8>, %arg2 : tensor<512xi32>) -> () {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 10, 10, 4, 128]
   // CHECK: [[CST0:%.+]] = arith.constant 0
   // CHECK: [[FILL:%.+]] = linalg.fill([[CST0]], [[INIT]])
   // CHECK: [[OUT:%.+]] = linalg.init_tensor [1, 10, 10, 512]
   // CHECK: [[C128:%.+]] = arith.constant -128
   // CHECK: [[C42:%.+]] = arith.constant 42
   // CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm_q {dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1, [[C128]], [[C42]] : tensor<1x14x14x4xi8>, tensor<3x3x4x128xi8>, i32, i32) outs([[FILL]] : tensor<1x10x10x4x128xi32>)
-  // CHECK: [[COLLAPSED:%.+]] = linalg.tensor_collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
+  // CHECK: [[COLLAPSED:%.+]] = tensor.collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
   // CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<512xi32>, tensor<1x10x10x512xi32>) outs([[OUT]] : tensor<1x10x10x512xi32>) {
   // CHECK: ^bb0(%arg3: i32, %arg4: i32, %arg5: i32):  // no predecessors
   // CHECK:   [[ADD:%.+]] = arith.addi %arg3, %arg4 : i32
   // CHECK:   linalg.yield [[ADD]] : i32
   // CHECK: } -> tensor<1x10x10x512xi32>
   %0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [0, 0, 0, 0], quantization_info = {input_zp = -128 : i32, weight_zp = 42 : i32}, stride = [1, 1], dilation = [2, 2] } : (tensor<1x14x14x4xi8>, tensor<3x3x4x128xi8>, tensor<512xi32>)  -> tensor<1x10x10x512xi32>
   return
 }

mlir/test/Dialect/Linalg/bufferize.mlir

func @bufferize_fill(%arg0: tensor<?xf32>) -> tensor<?xf32> {
   return %0 : tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @bufferize_tensor_collapse_shape(
 // CHECK-SAME:    %[[IN:.*]]: tensor<4x5xf32>
 func @bufferize_tensor_collapse_shape(%arg0: tensor<4x5xf32>) -> tensor<20xf32> {
-  %out = linalg.tensor_collapse_shape %arg0 [[0, 1]] :
+  %out = tensor.collapse_shape %arg0 [[0, 1]] :
      tensor<4x5xf32> into tensor<20xf32>
   return %out : tensor<20xf32>
 }
 // CHECK: %[[MEMREF:.*]] = bufferization.to_memref %[[IN]] : memref<4x5xf32>
 // CHECK: %[[RESHAPE:.*]] = memref.collapse_shape %[[MEMREF]] {{\[}}[0, 1]]
 // CHECK-SAME: : memref<4x5xf32> into memref<20xf32>
 // CHECK: %[[TENSOR:.*]] = bufferization.to_tensor %[[RESHAPE]] : memref<20xf32>
 // CHECK: return %[[TENSOR]]

mlir/test/Dialect/Linalg/canonicalize.mlir

linalg.tiled_loop (%arg3) = (%c0) to (%c192) step (%c24)
     linalg.fill(%cst, %16) : f32, memref<?xf32, #map>
     linalg.yield
   }
   return
 }
 
 // -----
 
-// CHECK-LABEL: zero_rank_reshape_multi
-func @zero_rank_reshape_multi(%arg0: tensor<f32>) -> tensor<f32> {
-  // CHECK: return %arg0
-  %0 = linalg.tensor_expand_shape %arg0 [] : tensor<f32> into tensor<1xf32>
-  %1 = linalg.tensor_expand_shape %0 [[0, 1]] : tensor<1xf32> into tensor<1x1xf32>
-  %2 = linalg.tensor_collapse_shape %1 [] : tensor<1x1xf32> into tensor<f32>
-  return %2 : tensor<f32>
-}
-
-// -----
-
-func @collapsing_tensor_reshapes(%arg0 : tensor<?x?x?x?x?xf32>) -> tensor<?x?xf32>
-{
-  %0 = linalg.tensor_collapse_shape %arg0 [[0, 1], [2], [3, 4]]
-      : tensor<?x?x?x?x?xf32> into tensor<?x?x?xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0, 1], [2]]
-      : tensor<?x?x?xf32> into tensor<?x?xf32>
-  return %1 : tensor<?x?xf32>
-}
-// CHECK-LABEL: collapsing_tensor_reshapes
-//       CHECK:   linalg.tensor_collapse_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
-//   CHECK-NOT:   linalg.tensor_collapse_shape
-
-// -----
-
-func @collapsing_tensor_reshapes_to_zero_dim(%arg0 : tensor<1x1x1xf32>)
-                                             -> tensor<f32> {
-  %0 = linalg.tensor_collapse_shape %arg0 [[0, 1, 2]]
-      : tensor<1x1x1xf32> into tensor<1xf32>
-  %1 = linalg.tensor_collapse_shape %0 [] : tensor<1xf32> into tensor<f32>
-  return %1 : tensor<f32>
-}
-// CHECK-LABEL: collapsing_tensor_reshapes_to_zero
-//       CHECK:   linalg.tensor_collapse_shape %{{.*}} []
-//  CHECK-SAME:     tensor<1x1x1xf32> into tensor<f32>
-
-// -----
-
-func @expanding_tensor_reshapes(%arg0 : tensor<?x?xf32>) -> tensor<?x6x4x?x5xf32>
-{
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1], [2]]
-      : tensor<?x?xf32> into tensor<?x4x?xf32>
-  %1 = linalg.tensor_expand_shape %0 [[0, 1], [2], [3, 4]]
-      : tensor<?x4x?xf32> into tensor<?x6x4x?x5xf32>
-  return %1 : tensor<?x6x4x?x5xf32>
-}
-// CHECK-LABEL: expanding_tensor_reshapes
-//       CHECK:   linalg.tensor_expand_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
-//   CHECK-NOT:   linalg.tensor_expand_shape
-
-// -----
-
-func @expanding_tensor_reshapes_to_zero_dim(%arg0 : tensor<f32>)
-                                             -> tensor<1x1x1xf32> {
-  %0 = linalg.tensor_expand_shape %arg0 [] : tensor<f32> into tensor<1xf32>
-  %1 = linalg.tensor_expand_shape %0 [[0, 1, 2]]
-      : tensor<1xf32> into tensor<1x1x1xf32>
-  return %1 : tensor<1x1x1xf32>
-}
-// CHECK-LABEL: expanding_tensor_reshapes_to_zero
-//       CHECK:   linalg.tensor_expand_shape %{{.*}} []
-//  CHECK-SAME:     tensor<f32> into tensor<1x1x1xf32>
-
-// -----
-
-func @fold_tensor_reshape(%arg0 : tensor<12x4xf32>) -> tensor<12x4xf32>
-{
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1], [2]]
-      : tensor<12x4xf32> into tensor<3x4x4xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0, 1], [2]]
-      : tensor<3x4x4xf32> into tensor<12x4xf32>
-  return %1 : tensor<12x4xf32>
-}
-// CHECK-LABEL: @fold_tensor_reshape
-//   CHECK-NOT:   linalg.{{.*}}shape
-
-// -----
-
-func @fold_tensor_reshape_dynamic(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>
-{
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1], [2]]
-      : tensor<?x?xf32> into tensor<?x4x?xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0, 1], [2]]
-      : tensor<?x4x?xf32> into tensor<?x?xf32>
-  return %1 : tensor<?x?xf32>
-}
-// CHECK-LABEL: @fold_tensor_reshape_dynamic
-//   CHECK-NOT:   linalg.{{.*}}_shape
-
-// -----
-
-func @reshape_collapse(%arg0 : tensor<2x3x4x5x6x7x8xf32>) -> tensor<24x5x42x8xf32>
-{
-  %0 = linalg.tensor_collapse_shape %arg0 [[0, 1, 2, 3, 4, 5, 6]]
-      : tensor<2x3x4x5x6x7x8xf32> into tensor<40320xf32>
-  %1 = linalg.tensor_expand_shape %0 [[0, 1, 2, 3]]
-      : tensor<40320xf32> into tensor<24x5x42x8xf32>
-  return %1 : tensor<24x5x42x8xf32>
-}
-//      CHECK: func @reshape_collapse
-// CHECK-SAME:   %[[ARG0:.+]]: tensor<2x3x4x5x6x7x8xf32>
-//      CHECK:   %[[RESULT:.+]] = linalg.tensor_collapse_shape %[[ARG0]]
-// CHECK-SAME:     [0, 1, 2], [3], [4, 5], [6]
-//      CHECK:   return %[[RESULT]]
-
-// -----
-
-func @reshape_expand(%arg0 : tensor<24x5x42x8xf32>) -> tensor<2x3x4x5x6x7x8xf32>
-{
-  %0 = linalg.tensor_collapse_shape %arg0 [[0, 1, 2, 3]]
-      : tensor<24x5x42x8xf32> into tensor<40320xf32>
-  %1 = linalg.tensor_expand_shape %0 [[0, 1, 2, 3, 4, 5, 6]]
-      : tensor<40320xf32> into tensor<2x3x4x5x6x7x8xf32>
-  return %1 : tensor<2x3x4x5x6x7x8xf32>
-}
-//      CHECK: func @reshape_expand
-// CHECK-SAME:   %[[ARG0:.+]]: tensor<24x5x42x8xf32>
-//      CHECK:   %[[RESULT:.+]] = linalg.tensor_expand_shape %[[ARG0]]
-// CHECK-SAME:     [0, 1, 2], [3], [4, 5], [6]
-//      CHECK:   return %[[RESULT]]
-
-// -----
-
-func @expand_reshape_1D(%arg0 : tensor<2048xf32>) -> tensor<4x512xf32>
-{
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1, 2, 3]]
-    : tensor<2048xf32> into tensor<1x4x1x512xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0, 1, 2], [3]]
-    : tensor<1x4x1x512xf32> into tensor<4x512xf32>
-  return %1 : tensor<4x512xf32>
-}
-//       CHECK: func @expand_reshape_1D
-//       CHECK: linalg.tensor_expand_shape %{{.*}} {{\[}}[0, 1]]
-//  CHECK-SAME:   tensor<2048xf32> into tensor<4x512xf32>
-
-// -----
-
-func @fold_reshape_1D(%arg0 : tensor<4x512xf32>) -> tensor<2048xf32>
-{
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1, 2], [3]]
-    : tensor<4x512xf32> into tensor<1x4x1x512xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0, 1, 2, 3]]
-    : tensor<1x4x1x512xf32> into tensor<2048xf32>
-  return %1 : tensor<2048xf32>
-}
-//       CHECK: func @fold_reshape_1D
-//       CHECK: linalg.tensor_collapse_shape %{{.*}} {{\[}}[0, 1]]
-//  CHECK-SAME:   tensor<4x512xf32> into tensor<2048xf32>
-
-// -----
-
-func @fold_reshape_unit_dims(%arg0 : tensor<2048x1x1xf32>) -> tensor<4x512x1x1xf32>
-{
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1, 2, 3], [4], [5]]
-    : tensor<2048x1x1xf32> into tensor<1x4x1x512x1x1xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0, 1, 2], [3], [4], [5]]
-    : tensor<1x4x1x512x1x1xf32> into tensor<4x512x1x1xf32>
-  return %1 : tensor<4x512x1x1xf32>
-}
-//       CHECK: func @fold_reshape_unit_dims
-//       CHECK: linalg.tensor_expand_shape %{{.*}} {{\[}}[0, 1], [2], [3]]
-//  CHECK-SAME:   tensor<2048x1x1xf32> into tensor<4x512x1x1xf32>
-
-// -----
-
-func @expand_reshape_unit_dims(%arg0 : tensor<2048x1x2048xf32>) -> tensor<4x512x1x512x4xf32>
-{
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1, 2, 3, 4], [5], [6, 7, 8]]
-    : tensor<2048x1x2048xf32> into tensor<1x4x1x512x1x1x512x1x4xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0, 1, 2], [3, 4], [5], [6, 7], [8]]
-    : tensor<1x4x1x512x1x1x512x1x4xf32> into tensor<4x512x1x512x4xf32>
-  return %1 : tensor<4x512x1x512x4xf32>
-}
-//       CHECK: func @expand_reshape_unit_dims
-//       CHECK: linalg.tensor_expand_shape %{{.*}} {{\[}}[0, 1], [2], [3, 4]]
-//  CHECK-SAME:   tensor<2048x1x2048xf32> into tensor<4x512x1x512x4xf32>
-
-// -----
-
-func @fold_reshape_trailing_unit_dims(%arg0: tensor<2xf32>) -> tensor<2x1xf32>
-{
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1, 2]]
-      : tensor<2xf32> into tensor<2x1x1xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0], [1, 2]]
-      : tensor<2x1x1xf32> into tensor<2x1xf32>
-  return %1 : tensor<2x1xf32>
-}
-//       CHECK: func @fold_reshape_trailing_unit_dims
-//       CHECK: linalg.tensor_expand_shape %{{.*}} {{\[}}[0, 1]]
-//  CHECK-SAME:   tensor<2xf32> into tensor<2x1xf32>
-
-// -----
-
-func @collapse_reshape_unit_dims_dynamic(%arg0 : tensor<?x1x?x1x1x?x?x1x1xf32>) -> tensor<?x?x?x?xf32>
-{
-  %0 = linalg.tensor_collapse_shape %arg0 [[0], [1, 2], [3], [4], [5], [6, 7, 8]]
-    : tensor<?x1x?x1x1x?x?x1x1xf32> into tensor<?x?x1x1x?x?xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0], [1], [2, 3, 4], [5]]
-    : tensor<?x?x1x1x?x?xf32> into tensor<?x?x?x?xf32>
-  return %1 : tensor<?x?x?x?xf32>
-}
-//       CHECK: func @collapse_reshape_unit_dims_dynamic
-//       CHECK: linalg.tensor_collapse_shape
-//  CHECK-SAME:   [0], [1, 2], [3, 4, 5], [6, 7, 8]
-//  CHECK-SAME:   tensor<?x1x?x1x1x?x?x1x1xf32> into tensor<?x?x?x?xf32>
-
-// -----
-
-func @fold_reshape_trailing_unit_dims(%arg0: tensor<2xf32>) -> tensor<2x1xf32>
-{
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1, 2]]
-      : tensor<2xf32> into tensor<2x1x1xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0], [1, 2]]
-      : tensor<2x1x1xf32> into tensor<2x1xf32>
-  return %1 : tensor<2x1xf32>
-}
-//       CHECK: func @fold_reshape_trailing_unit_dims
-//       CHECK: linalg.tensor_expand_shape %{{.*}} {{\[}}[0, 1]]
-//  CHECK-SAME:   tensor<2xf32> into tensor<2x1xf32>
-
-// -----
-
-func @fold_reshape_trailing_unit_dims_dynamic(%arg0: tensor<1x1x?x1x1x1xf32>) -> tensor<?xf32>
-{
-  %0 = linalg.tensor_collapse_shape %arg0 [[0, 1, 2], [3], [4], [5]]
-      : tensor<1x1x?x1x1x1xf32> into tensor<?x1x1x1xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0, 1, 2, 3]]
-      : tensor<?x1x1x1xf32> into tensor<?xf32>
-  return %1 : tensor<?xf32>
-}
-//       CHECK: func @fold_reshape_trailing_unit_dims_dynamic
-//       CHECK: linalg.tensor_collapse_shape %{{.*}} {{\[}}[0, 1, 2, 3, 4, 5]]
-//  CHECK-SAME:   tensor<1x1x?x1x1x1xf32> into tensor<?xf32>
-
-// -----
-
-func @fold_reshape_trailing_unit_dims(%arg0: tensor<12x42x1x1xf32>) -> tensor<12x42xf32>
-{
-  %0 = linalg.tensor_expand_shape %arg0 [[0], [1], [2], [3, 4]]
-      : tensor<12x42x1x1xf32> into tensor<12x42x1x1x1xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0], [1, 2, 3, 4]]
-      : tensor<12x42x1x1x1xf32> into tensor<12x42xf32>
-  return %1 : tensor<12x42xf32>
-}
-//       CHECK: func @fold_reshape_trailing_unit_dims
-//       CHECK: linalg.tensor_collapse_shape %{{.*}} {{\[}}[0], [1, 2, 3]]
-//  CHECK-SAME:   tensor<12x42x1x1xf32> into tensor<12x42xf32>
-
-// -----
-
-func @no_fold_reshapes(%arg0 : tensor<?x?x?xf32>) -> tensor<?x?xf32>
-{
-  %0 = linalg.tensor_expand_shape %arg0 [[0], [1], [2, 3]]
-      : tensor<?x?x?xf32> into tensor<?x?x1x?xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0], [1, 2, 3]]
-      : tensor<?x?x1x?xf32> into tensor<?x?xf32>
-  return %1 : tensor<?x?xf32>
-}
-// CHECK-LABEL: func @no_fold_reshapes
-//       CHECK:   linalg.tensor_expand_shape
-//       CHECK:   linalg.tensor_collapse_shape
-
-// -----
-
-func @no_fold_reshape_incompatible(%arg0 : tensor<4x6x8xf32>) -> tensor<2x6x16xf32>
-{
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1], [2, 3], [4]]
-      : tensor<4x6x8xf32> into tensor<2x2x3x2x8xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0], [1, 2], [3, 4]]
-      : tensor<2x2x3x2x8xf32> into tensor<2x6x16xf32>
-  return %1 : tensor<2x6x16xf32>
-}
-// CHECK-LABEL: func @no_fold_reshape_incompatible
-//       CHECK:   linalg.tensor_expand_shape
-//       CHECK:   linalg.tensor_collapse_shape
-
-// -----
-
-func @no_fold_reshape_empty_expr(%arg0: tensor<3x2x2xf32>) -> tensor<12x1xf32> {
-  %0 = linalg.tensor_expand_shape %arg0 [[0], [1], [2, 3]]
-      : tensor<3x2x2xf32> into tensor<3x2x2x1xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0, 1, 2], [3]]
-      : tensor<3x2x2x1xf32> into tensor<12x1xf32>
-  return %1 : tensor<12x1xf32>
-}
-//      CHECK: func @no_fold_reshape_empty_expr
-// CHECK-SAME:    %[[ARG0:.+]]: tensor<3x2x2xf32>
-//      CHECK:    %[[RARG0:.+]] = linalg.tensor_expand_shape %[[ARG0]]
-// CHECK-SAME:      [0], [1], [2, 3]
-//      CHECK:    %[[RES:.+]] = linalg.tensor_collapse_shape %[[RARG0]]
-// CHECK-SAME:      [0, 1, 2], [3]
-//      CHECK:    return %[[RES:.+]] : tensor<12x1xf32>
-
-// -----
-
 #accesses = [
   affine_map<(i) -> (i)>
 ]
 
 #trait = {
   indexing_maps = #accesses,
   iterator_types = ["parallel"]
 }

 // CHECK-LABEL: @dce_zero_memref
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: memref<0xf32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<0xf32>
 //   CHECK-NOT:   linalg.copy
 //  CHECK-NEXT:   return %[[ARG1]]
 
 // -----
 
-func @reshape_splat_constant_int32() -> tensor<2x4x2xi32>
-{
-  %c0 = arith.constant dense<42> : tensor<2x8xi32>
-  %0 = linalg.tensor_expand_shape %c0 [[0], [1, 2]]
-      : tensor<2x8xi32> into tensor<2x4x2xi32>
-  return %0 : tensor<2x4x2xi32>
-}
-// CHECK-LABEL: @reshape_splat_constant_int32
-//       CHECK:   %[[CST:.*]] = arith.constant dense<{{.*}}> : tensor<2x4x2xi32>
-//   CHECK-NOT:   linalg.tensor_expand_shape
-//       CHECK:   return %[[CST]]
-
-func @reshape_splat_constant_int16() -> tensor<2x4x2xi16>
-{
-  %c0 = arith.constant dense<42> : tensor<2x8xi16>
-  %0 = linalg.tensor_expand_shape %c0 [[0], [1, 2]]
-      : tensor<2x8xi16> into tensor<2x4x2xi16>
-  return %0 : tensor<2x4x2xi16>
-}
-// CHECK-LABEL: @reshape_splat_constant_int16
-//       CHECK:   %[[CST:.*]] = arith.constant dense<{{.*}}> : tensor<2x4x2xi16>
-//   CHECK-NOT:   linalg.tensor_expand_shape
-//       CHECK:   return %[[CST]]
-
-func @reshape_splat_constant_float32() -> tensor<2x4x2xf32>
-{
-  %c0 = arith.constant dense<42.0> : tensor<2x8xf32>
-  %0 = linalg.tensor_expand_shape %c0 [[0], [1, 2]]
-      : tensor<2x8xf32> into tensor<2x4x2xf32>
-  return %0 : tensor<2x4x2xf32>
-}
-// CHECK-LABEL: @reshape_splat_constant_float32
-//       CHECK:   %[[CST:.*]] = arith.constant dense<{{.*}}> : tensor<2x4x2xf32>
-//   CHECK-NOT:   linalg.tensor_expand_shape
-//       CHECK:   return %[[CST]]
-
-func @reshape_splat_constant_float64() -> tensor<2x4x2xf64>
-{
-  %c0 = arith.constant dense<42.0> : tensor<2x8xf64>
-  %0 = linalg.tensor_expand_shape %c0 [[0], [1, 2]]
-      : tensor<2x8xf64> into tensor<2x4x2xf64>
-  return %0 : tensor<2x4x2xf64>
-}
-// CHECK-LABEL: @reshape_splat_constant_float64
-//       CHECK:   %[[CST:.*]] = arith.constant dense<{{.*}}> : tensor<2x4x2xf64>
-//   CHECK-NOT:   linalg.tensor_expand_shape
-//       CHECK:   return %[[CST]]
-
-// -----
 
 // CHECK-LABEL: func @tensor.cast(
 func @tensor.cast(%a : tensor<3x4xf32>, %b : tensor<4x?xf32>, %c : tensor<3x?xf32>)
   -> tensor<3x?xf32>
 {
   %ta = tensor.cast %a : tensor<3x4xf32> to tensor<?x?xf32>
   %tb = tensor.cast %b : tensor<4x?xf32> to tensor<?x?xf32>
   %tc = tensor.cast %c : tensor<3x?xf32> to tensor<?x?xf32>

 // CHECK:   %[[T0:.+]] = linalg.init_tensor [4, 5, 6] : tensor<4x5x6xf32>
 // CHECK:   %[[T1:.+]] = tensor.cast %[[T0]] : tensor<4x5x6xf32> to tensor<4x5x?xf32>
 // CHECK:   return %[[T1]]
 
 // -----
 
 func @init_tensor_reshape_expansion(%arg0 : index) -> tensor<2x3x5x4x?x7xf32> {
   %0 = linalg.init_tensor [6, 5, %arg0] : tensor<6x5x?xf32>
-  %1 = linalg.tensor_expand_shape %0 [[0, 1], [2], [3, 4, 5]]
+  %1 = tensor.expand_shape %0 [[0, 1], [2], [3, 4, 5]]
       : tensor<6x5x?xf32> into tensor<2x3x5x4x?x7xf32>
   return %1 : tensor<2x3x5x4x?x7xf32>
 }
 //      CHECK: #[[MAP:.+]] = affine_map<()[s0] -> (s0 floordiv 28)>
 //      CHECK: func @init_tensor_reshape_expansion
 // CHECK-SAME:     %[[ARG0:.+]]: index
 // CHECK-NEXT:   %[[D:.+]] = affine.apply #[[MAP]]()[%[[ARG0]]]
 // CHECK-NEXT:   %[[INIT:.+]] = linalg.init_tensor [2, 3, 5, 4, %[[D]], 7]
 // CHECK-NEXT:   return %[[INIT]]
 
 // -----
 
 func @init_tensor_reshape_collapse(%arg0 : index) -> tensor<6x5x?xf32> {
   %0 = linalg.init_tensor [2, 3, 5, 4, %arg0, 7] : tensor<2x3x5x4x?x7xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0, 1], [2], [3, 4, 5]]
+  %1 = tensor.collapse_shape %0 [[0, 1], [2], [3, 4, 5]]
       : tensor<2x3x5x4x?x7xf32> into tensor<6x5x?xf32>
   return %1 : tensor<6x5x?xf32>
 }
 //      CHECK: #[[MAP:.+]] = affine_map<()[s0] -> (s0 * 28)>
 //      CHECK: func @init_tensor_reshape_collapse
 // CHECK-SAME:     %[[ARG0:.+]]: index
 // CHECK-NEXT:   %[[D:.+]] = affine.apply #[[MAP]]()[%[[ARG0]]]
 // CHECK-NEXT:   %[[INIT:.+]] = linalg.init_tensor [6, 5, %[[D]]]

 
 // CHECK-LABEL: func @fold_fill_reshape()
 func @fold_fill_reshape() -> tensor<6x4xf32> {
   %zero = arith.constant 0.0 : f32
   // CHECK: %[[INIT:.+]] = linalg.init_tensor [6, 4] : tensor<6x4xf32>
   %init = linalg.init_tensor [1, 2, 3, 4] : tensor<1x2x3x4xf32>
   // CHECK: %[[FILL:.+]] = linalg.fill(%cst, %[[INIT]]) : f32, tensor<6x4xf32> -> tensor<6x4xf32>
   %fill = linalg.fill(%zero, %init) : f32, tensor<1x2x3x4xf32> -> tensor<1x2x3x4xf32>
-  %reshape = linalg.tensor_collapse_shape %fill [[0, 1, 2], [3]]
+  %reshape = tensor.collapse_shape %fill [[0, 1, 2], [3]]
       : tensor<1x2x3x4xf32> into tensor<6x4xf32>
   // CHECK: return %[[FILL]] : tensor<6x4xf32>
   return %reshape : tensor<6x4xf32>
 }
 
 // -----
 
 //       CHECK: func @fold_fill_reshape_dynamic
 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<?x?x?x?x?xf32>
 func @fold_fill_reshape_dynamic(%arg0 : tensor<?x?x?x?x?xf32>) -> tensor<?x?xf32> {
   %zero = arith.constant 0.0 : f32
-  // CHECK: %[[RESHAPE:.+]] = linalg.tensor_collapse_shape %[[ARG0]]
+  // CHECK: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]]
   %0 = linalg.fill(%zero, %arg0) : f32, tensor<?x?x?x?x?xf32> -> tensor<?x?x?x?x?xf32>
   // CHECK: %[[RESULT:.+]] = linalg.fill(%{{.+}}, %[[RESHAPE]])
-  %1 = linalg.tensor_collapse_shape %0 [[0, 1, 2], [3, 4]]
+  %1 = tensor.collapse_shape %0 [[0, 1, 2], [3, 4]]
       : tensor<?x?x?x?x?xf32> into tensor<?x?xf32>
   // CHECK: return %[[RESULT]]
   return %1 : tensor<?x?xf32>
 }
 
 
 // -----
 
func @dim_of_tiled_loop_result_no_canonicalize(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>, %s: index)
   %r2 = tensor.dim %r, %c0 : tensor<?x?xf32>
   return %r2 : index
 }
 
 // -----
 
 // CHECK-LABEL: @depthwise_conv
 func @depthwise_conv(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<?x?x?x1xf32>, %arg2: tensor<?x?x?x?x1xf32>) -> tensor<?x?x?x?x1xf32> {
-  // CHECK-DAG: %[[KERNEL:.+]] = linalg.tensor_collapse_shape %arg1 {{\[\[}}0], [1], [2, 3]]
-  // CHECK-DAG: %[[INIT:.+]] = linalg.tensor_collapse_shape %arg2 {{\[\[}}0], [1], [2], [3, 4]]
+  // CHECK-DAG: %[[KERNEL:.+]] = tensor.collapse_shape %arg1 {{\[\[}}0], [1], [2, 3]]
+  // CHECK-DAG: %[[INIT:.+]] = tensor.collapse_shape %arg2 {{\[\[}}0], [1], [2], [3, 4]]
   // CHECK-DAG: %[[CONV:.+]] = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %[[KERNEL]] : tensor<?x?x?x?xf32>, tensor<?x?x?xf32>) outs(%[[INIT]] : tensor<?x?x?x?xf32>)
-  // CHECK: %[[OUT:.+]] = linalg.tensor_expand_shape %[[CONV]] {{\[\[}}0], [1], [2], [3, 4]]
+  // CHECK: %[[OUT:.+]] = tensor.expand_shape %[[CONV]] {{\[\[}}0], [1], [2], [3, 4]]
   %0 = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<?x?x?x?xf32>, tensor<?x?x?x1xf32>) outs(%arg2 : tensor<?x?x?x?x1xf32>) -> tensor<?x?x?x?x1xf32>
   return %0 : tensor<?x?x?x?x1xf32>
 }
 
 
 // -----
 
 // CHECK-LABEL: @depthwise_conv_q
 func @depthwise_conv_q(%arg0: tensor<?x?x?x?xi8>, %arg1: tensor<?x?x?x1xi8>, %arg2: tensor<?x?x?x?x1xi32>, %arg3 : i32, %arg4 : i32) -> tensor<?x?x?x?x1xi32> {
-  // CHECK-DAG: %[[KERNEL:.+]] = linalg.tensor_collapse_shape %arg1 {{\[\[}}0], [1], [2, 3]]
-  // CHECK-DAG: %[[INIT:.+]] = linalg.tensor_collapse_shape %arg2 {{\[\[}}0], [1], [2], [3, 4]]
+  // CHECK-DAG: %[[KERNEL:.+]] = tensor.collapse_shape %arg1 {{\[\[}}0], [1], [2, 3]]
+  // CHECK-DAG: %[[INIT:.+]] = tensor.collapse_shape %arg2 {{\[\[}}0], [1], [2], [3, 4]]
   // CHECK-DAG: %[[CONV:.+]] = linalg.depthwise_conv_2d_nhwc_hwc_q {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %[[KERNEL]], %arg3, %arg4 : tensor<?x?x?x?xi8>, tensor<?x?x?xi8>, i32, i32) outs(%[[INIT]] : tensor<?x?x?x?xi32>)
-  // CHECK: %[[OUT:.+]] = linalg.tensor_expand_shape %[[CONV]] {{\[\[}}0], [1], [2], [3, 4]]
+  // CHECK: %[[OUT:.+]] = tensor.expand_shape %[[CONV]] {{\[\[}}0], [1], [2], [3, 4]]
   %0 = linalg.depthwise_conv_2d_nhwc_hwcm_q {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %arg1, %arg3, %arg4 : tensor<?x?x?x?xi8>, tensor<?x?x?x1xi8>, i32, i32) outs(%arg2 : tensor<?x?x?x?x1xi32>) -> tensor<?x?x?x?x1xi32>
   return %0 : tensor<?x?x?x?x1xi32>
 }

mlir/test/Dialect/Linalg/detensorize_0d.mlir

func @detensor_simple(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} {
   return %1: tensor<f32>
 }
 // CHECK-LABEL: func @detensor_simple
 // CHECK-SAME:    (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>)
 // CHECK-DAG:     %[[arg1_val:.*]] = tensor.extract %[[arg1]]
 // CHECK-DAG:     %[[arg2_val:.*]] = tensor.extract %[[arg2]]
 // CHECK:         %[[detensored_res:.*]] = arith.addf %[[arg1_val]], %[[arg2_val]]
 // CHECK:         %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res]]
-// CHECK:         %[[reshaped_tensor_res:.*]] = linalg.tensor_collapse_shape %[[new_tensor_res]]
+// CHECK:         %[[reshaped_tensor_res:.*]] = tensor.collapse_shape %[[new_tensor_res]]
 // CHECK:         return %[[reshaped_tensor_res]]
 
 func @detensor_op_sequence(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} {
   %0 = linalg.init_tensor [] : tensor<f32>
   %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []}
     ins(%arg1, %arg2 : tensor<f32>, tensor<f32>)
     outs(%0 : tensor<f32>) {
   ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):  // no predecessors

 // CHECK-LABEL: func @detensor_op_sequence
 // CHECK-SAME:    (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>)
 // CHECK-DAG:     %[[arg1_val:.*]] = tensor.extract %[[arg1]]
 // CHECK-DAG:     %[[arg2_val:.*]] = tensor.extract %[[arg2]]
 // CHECK:         %[[detensored_res:.*]] = arith.addf %[[arg1_val]], %[[arg2_val]]
 // CHECK:         %[[detensored_res2:.*]] = arith.mulf %[[arg1_val]], %[[detensored_res]]
 // CHECK:         %[[detensored_res3:.*]] = arith.divf %[[detensored_res]], %[[detensored_res2]]
 // CHECK:         %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res3]]
-// CHECK:         %[[reshaped_tensor_res:.*]] = linalg.tensor_collapse_shape %[[new_tensor_res]]
+// CHECK:         %[[reshaped_tensor_res:.*]] = tensor.collapse_shape %[[new_tensor_res]]
 // CHECK:         return %[[reshaped_tensor_res]]
 
 func @detensor_multiple_ops(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} {
   %0 = linalg.init_tensor [] : tensor<f32>
   %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []}
     ins(%arg1, %arg2 : tensor<f32>, tensor<f32>)
     outs(%0 : tensor<f32>) {
   ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):  // no predecessors
     %2 = arith.addf %arg3, %arg4 : f32
     %3 = arith.mulf %2, %arg4 : f32
     linalg.yield %3 : f32
   } -> tensor<f32>
   return %1: tensor<f32>
 }
 // CHECK-LABEL: func @detensor_multiple_ops
 // CHECK-SAME:    (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>)
 // CHECK-DAG:     %[[arg1_val:.*]] = tensor.extract %[[arg1]]
 // CHECK-DAG:     %[[arg2_val:.*]] = tensor.extract %[[arg2]]
 // CHECK:         %[[detensored_res:.*]] = arith.addf %[[arg1_val]], %[[arg2_val]]
 // CHECK:         %[[detensored_res2:.*]] = arith.mulf %[[detensored_res]], %[[arg2_val]]
 // CHECK:         %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res2]]
-// CHECK:         %[[reshaped_tensor_res:.*]] = linalg.tensor_collapse_shape %[[new_tensor_res]]
+// CHECK:         %[[reshaped_tensor_res:.*]] = tensor.collapse_shape %[[new_tensor_res]]
 // CHECK:         return %[[reshaped_tensor_res]]
 
 func @detensor_foreign_op(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} {
   %0 = linalg.init_tensor [] : tensor<f32>
   %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []}
     ins(%arg1, %arg2 : tensor<f32>, tensor<f32>)
     outs(%0 : tensor<f32>) {
   ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):  // no predecessors
     %2 = "foreign.do_something"(%arg3, %arg4) {} : (f32, f32) -> f32
     linalg.yield %2 : f32
   } -> tensor<f32>
   return %1: tensor<f32>
 }
 // CHECK-LABEL: func @detensor_foreign_op
 // CHECK-SAME:    (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>)
 // CHECK-DAG:     %[[arg1_val:.*]] = tensor.extract %[[arg1]]
 // CHECK-DAG:     %[[arg2_val:.*]] = tensor.extract %[[arg2]]
 // CHECK:         %[[detensored_res:.*]] = "foreign.do_something"(%[[arg1_val]], %[[arg2_val]])
 // CHECK:         %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res]]
-// CHECK:         %[[reshaped_tensor_res:.*]] = linalg.tensor_collapse_shape %[[new_tensor_res]]
+// CHECK:         %[[reshaped_tensor_res:.*]] = tensor.collapse_shape %[[new_tensor_res]]
 // CHECK:         return %[[reshaped_tensor_res]]

mlir/test/Dialect/Linalg/detensorize_br_operands.mlir

 // RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -linalg-detensorize | FileCheck %s
 
 // TODO: Detensoring breaks if %arg0 or %arg1 are passed directly as tensors. Fix that.
 func @if_true_test(%arg0: i1, %arg1: i32) -> tensor<i32> attributes {} {
   %arg0_t = tensor.from_elements %arg0 : tensor<1xi1>
-  %arg0_t2 = linalg.tensor_collapse_shape %arg0_t [] : tensor<1xi1> into tensor<i1>
+  %arg0_t2 = tensor.collapse_shape %arg0_t [] : tensor<1xi1> into tensor<i1>
 
   %arg1_t = tensor.from_elements %arg1 : tensor<1xi32>
-  %arg1_t2 = linalg.tensor_collapse_shape %arg1_t [] : tensor<1xi32> into tensor<i32>
+  %arg1_t2 = tensor.collapse_shape %arg1_t [] : tensor<1xi32> into tensor<i32>
 
   %cst = arith.constant dense<10> : tensor<i32>
   %2 = linalg.init_tensor [] : tensor<i8>
   %3 = linalg.generic
     {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []}
     ins(%arg0_t2 : tensor<i1>)
     outs(%2 : tensor<i8>) {
   ^bb0(%arg2: i1, %arg3: i8):  // no predecessors

 // CHECK-SAME:     (%[[arg0:.*]]: i1, %[[arg1:.*]]: i32)
 // CHECK-NEXT:     arith.constant 10 : i32
 // CHECK-NEXT:     cond_br %[[arg0]], ^[[bb1:.*]], ^[[bb2:.*]](%[[arg1]] : i32)
 // CHECK-NEXT:   ^[[bb1]]:
 // CHECK-NEXT:     %[[add_res:.*]] = arith.addi
 // CHECK-NEXT:     br ^[[bb2]](%[[add_res]] : i32)
 // CHECK-NEXT:   ^[[bb2]]
 // CHECK-NEXT:     tensor.from_elements
-// CHECK-NEXT:     %[[func_res:.*]] = linalg.tensor_collapse_shape
+// CHECK-NEXT:     %[[func_res:.*]] = tensor.collapse_shape
 // CHECK-NEXT:     return %[[func_res]]

mlir/test/Dialect/Linalg/detensorize_if.mlir

 // RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -linalg-detensorize | FileCheck %s
 
 #map0 = affine_map<() -> ()>
 
 #attrs = {
   indexing_maps = [#map0, #map0, #map0],
   iterator_types = []
 }
 
 func @main() -> (tensor<i32>) attributes {} {
   %c0 = arith.constant 0 : i32
   %0 = tensor.from_elements %c0 : tensor<1xi32>
-  %reshaped0 = linalg.tensor_collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
+  %reshaped0 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
   %c10 = arith.constant 10 : i32
   %1 = tensor.from_elements %c10 : tensor<1xi32>
-  %reshaped1 = linalg.tensor_collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
+  %reshaped1 = tensor.collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
   br ^bb1(%reshaped0 : tensor<i32>)
 
 ^bb1(%2: tensor<i32>):  // 2 preds: ^bb0, ^bb2
   %3 = linalg.init_tensor [] : tensor<i1>
   %4 = linalg.generic #attrs
     ins(%2, %reshaped1 : tensor<i32>, tensor<i32>)
     outs(%3 : tensor<i1>) {
     ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):  // no predecessors

 // CHECK-NEXT:   ^[[bb1]](%{{.*}}: i32):
 // CHECK-NEXT:     arith.cmpi slt, %{{.*}}, %{{.*}}
 // CHECK-NEXT:     cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^bb3(%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb2]](%{{.*}}: i32)
 // CHECK-NEXT:     arith.addi %{{.*}}, %{{.*}}
 // CHECK-NEXT:     br ^[[bb3:.*]](%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb3]](%{{.*}}: i32)
 // CHECK-NEXT:     tensor.from_elements %{{.*}} : tensor<1xi32>
-// CHECK-NEXT:     linalg.tensor_collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
+// CHECK-NEXT:     tensor.collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
 // CHECK-NEXT:     return %{{.*}}
 // CHECK-NEXT:   }
 
 // -----
 
 // Similar to the above test with one change: one of the block after the
 // if-condition passes/forwards its tensor argument to another block.
 
 #map0 = affine_map<() -> ()>
 
 #attrs = {
   indexing_maps = [#map0, #map0, #map0],
   iterator_types = []
 }
 
 func @main() -> (tensor<i32>) attributes {} {
   %c0 = arith.constant 0 : i32
   %0 = tensor.from_elements %c0 : tensor<1xi32>
-  %reshaped0 = linalg.tensor_collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
+  %reshaped0 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
   %c10 = arith.constant 10 : i32
   %1 = tensor.from_elements %c10 : tensor<1xi32>
-  %reshaped1 = linalg.tensor_collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
+  %reshaped1 = tensor.collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
   br ^bb1(%reshaped0 : tensor<i32>)
 
 ^bb1(%2: tensor<i32>):  // 2 preds: ^bb0, ^bb2
   %3 = linalg.init_tensor [] : tensor<i1>
   %4 = linalg.generic #attrs
     ins(%2, %reshaped1 : tensor<i32>, tensor<i32>)
     outs(%3 : tensor<i1>) {
     ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):  // no predecessors

 // CHECK-NEXT:     cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^bb3(%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb2]](%{{.*}}: i32)
 // CHECK-NEXT:     arith.addi %{{.*}}, %{{.*}}
 // CHECK-NEXT:     br ^[[bb3:.*]](%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb3]](%{{.*}}: i32)
 // CHECK-NEXT:     br ^[[bb4:.*]](%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb4]](%{{.*}}: i32)
 // CHECK-NEXT:     tensor.from_elements %{{.*}} : tensor<1xi32>
-// CHECK-NEXT:     linalg.tensor_collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
+// CHECK-NEXT:     tensor.collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
 // CHECK-NEXT:     return %{{.*}}
 // CHECK-NEXT:   }
 
 // -----
 
 #map0 = affine_map<() -> ()>
 
 #attrs = {
   indexing_maps = [#map0, #map0, #map0],
   iterator_types = []
 }
 
 func @main() -> (tensor<i32>) attributes {} {
   %c0 = arith.constant 0 : i32
   %0 = tensor.from_elements %c0 : tensor<1xi32>
-  %reshaped0 = linalg.tensor_collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
+  %reshaped0 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
   %c10 = arith.constant 10 : i32
   %1 = tensor.from_elements %c10 : tensor<1xi32>
-  %reshaped1 = linalg.tensor_collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
+  %reshaped1 = tensor.collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
   br ^bb1(%reshaped0 : tensor<i32>)
 
 ^bb1(%2: tensor<i32>):  // 2 preds: ^bb0, ^bb2
   %3 = linalg.init_tensor [] : tensor<i1>
   %4 = linalg.generic #attrs
     ins(%2, %reshaped1 : tensor<i32>, tensor<i32>)
     outs(%3 : tensor<i1>) {
     ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):  // no predecessors
       %8 = arith.cmpi slt, %arg0, %arg1 : i32
       linalg.yield %8 : i1
   } -> tensor<i1>
   %5 = tensor.extract %4[] : tensor<i1>
   // This cond_br intentionally has bb2 as it's target for both branches. This
   // is to make sure that the "forward phase" of the cost-model correctly adds
   // the users of a block argument (in this case bb2's argument) to the work
   // list.
   cond_br %5, ^bb2(%2 : tensor<i32>), ^bb2(%2 : tensor<i32>)
 
 ^bb2(%6: tensor<i32>):  // pred: ^bb1
   %12 = tensor.from_elements %c10 : tensor<1xi32>
-  %reshaped12 = linalg.tensor_collapse_shape %12 [] : tensor<1xi32> into tensor<i32>
+  %reshaped12 = tensor.collapse_shape %12 [] : tensor<1xi32> into tensor<i32>
   %7 = linalg.init_tensor [] : tensor<i32>
   %8 = linalg.generic #attrs
     ins(%6, %reshaped12 : tensor<i32>, tensor<i32>)
     outs(%7 : tensor<i32>) {
     ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):  // no predecessors
       %9 = arith.addi %arg0, %arg1 : i32
       linalg.yield %9 : i32
   } -> tensor<i32>

 // CHECK-NEXT:   ^[[bb1]](%{{.*}}: i32):
 // CHECK-NEXT:     arith.cmpi slt, %{{.*}}, %{{.*}}
 // CHECK-NEXT:     cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^bb2(%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb2]](%{{.*}}: i32)
 // CHECK-NEXT:     arith.addi %{{.*}}, %{{.*}}
 // CHECK-NEXT:     br ^[[bb3:.*]](%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb3]](%{{.*}}: i32)
 // CHECK-NEXT:     tensor.from_elements %{{.*}} : tensor<1xi32>
-// CHECK-NEXT:     linalg.tensor_collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
+// CHECK-NEXT:     tensor.collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
 // CHECK-NEXT:     return %{{.*}}
 // CHECK-NEXT:   }

mlir/test/Dialect/Linalg/detensorize_trivial.mlir

 // RUN: mlir-opt %s -linalg-detensorize=aggressive-mode | FileCheck %s -check-prefix=DET-ALL
 // RUN: mlir-opt %s -linalg-detensorize | FileCheck %s -check-prefix=DET-CF
 
 
 #map0 = affine_map<() -> ()>
 
 #attrs = {
   indexing_maps = [#map0, #map0, #map0],
   iterator_types = []
 }
 
 func @main(%farg0 : tensor<i32>) -> (tensor<i1>) attributes {} {
   %c10 = arith.constant 10 : i32
   %1 = tensor.from_elements %c10 : tensor<1xi32>
-  %reshaped1 = linalg.tensor_collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
+  %reshaped1 = tensor.collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
   %3 = linalg.init_tensor [] : tensor<i1>
   %4 = linalg.generic #attrs
     ins(%farg0, %reshaped1 : tensor<i32>, tensor<i32>)
     outs(%3 : tensor<i1>) {
     ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
       %8 = arith.cmpi slt, %arg0, %arg1 : i32
       linalg.yield %8 : i1
   } -> tensor<i1>
   return %4 : tensor<i1>
 }
 
 
 // DET-ALL-LABEL: func @main(%{{.*}}: tensor<i32>)
 // DET-ALL-NEXT:    arith.constant 10
 // DET-ALL-NEXT:    tensor.extract %{{.*}}[]
 // DET-ALL-NEXT:    arith.cmpi slt, %{{.*}}, %{{.*}}
 // DET-ALL-NEXT:    tensor.from_elements %{{.*}}
-// DET-ALL-NEXT:    linalg.tensor_collapse_shape %{{.*}}
+// DET-ALL-NEXT:    tensor.collapse_shape %{{.*}}
 // DET-ALL-NEXT:    return %{{.*}} : tensor<i1>
 // DET-ALL-NEXT:  }
 
 // DET-CF-LABEL: func @main(%{{.*}}: tensor<i32>)
 // DET-CF-NEXT:    arith.constant dense<10> : tensor<i32>
 // DET-CF-NEXT:    linalg.init_tensor [] : tensor<i1>
 // DET-CF-NEXT:    linalg.generic
 // DET-CF-NEXT:    ^{{.*}}(%{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i1)
 // DET-CF-NEXT:      arith.cmpi slt, %{{.*}}, %{{.*}}
 // DET-CF-NEXT:      linalg.yield %{{.*}}
 // DET-CF-NEXT:    } -> tensor<i1>
 // DET-CF-NEXT:    return %{{.*}}
 // DET-CF-NEXT:  }

mlir/test/Dialect/Linalg/detensorize_while.mlir

 // DET-ALL:       ^[[bb1]](%{{.*}}: i32)
 // DET-ALL:         arith.cmpi slt, {{.*}}
 // DET-ALL:         cond_br {{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^[[bb3:.*]](%{{.*}} : i32)
 // DET-ALL:       ^[[bb2]](%{{.*}}: i32)
 // DET-ALL:         arith.addi {{.*}}
 // DET-ALL:         br ^[[bb1]](%{{.*}} : i32)
 // DET-ALL:       ^[[bb3]](%{{.*}}: i32)
 // DET-ALL:         tensor.from_elements {{.*}}
-// DET-ALL:         linalg.tensor_collapse_shape {{.*}}
+// DET-ALL:         tensor.collapse_shape {{.*}}
 // DET-ALL:         return %{{.*}} : tensor<i32>
 
 // Test detensoring only ops involed in control-flow.
 //
 // DET-CF-LABEL: func @main
 // DET-CF-SAME:    (%{{.*}}: tensor<i32>, %{{.*}}: tensor<i32>)
 // DET-CF:         tensor.extract {{.*}}
 // DET-CF:         br ^[[bb1:.*]](%{{.*}} : i32)
 // DET-CF:       ^[[bb1]](%{{.*}}: i32)
 // DET-CF:         arith.cmpi slt, {{.*}}
 // DET-CF:         cond_br {{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^[[bb3:.*]](%{{.*}} : i32)
 // DET-CF:       ^[[bb2]](%{{.*}}: i32)
 // DET-CF:         arith.addi {{.*}}
 // DET-CF:         br ^[[bb1]](%{{.*}} : i32)
 // DET-CF:       ^[[bb3]](%{{.*}}: i32)
 // DET-CF:         tensor.from_elements %{{.*}} : tensor<1xi32>
-// DET-CF:         linalg.tensor_collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
+// DET-CF:         tensor.collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
 // DET-CF:         return %{{.*}} : tensor<i32>

mlir/test/Dialect/Linalg/detensorize_while_impure_cf.mlir

 // DET-ALL:           %{{.*}} = arith.addi %{{.*}}, %{{.*}}
 // DET-ALL:           linalg.yield %{{.*}} : i32
 // DET-ALL:         } -> tensor<i32>
 // DET-ALL:         tensor.extract %{{.*}}[] : tensor<i32>
 // DET-ALL:         cmpi slt, %{{.*}}, %{{.*}} : i32
 // DET-ALL:         cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^[[bb3:.*]](%{{.*}} : i32)
 // DET-ALL:       ^[[bb2]](%{{.*}}: i32)
 // DET-ALL:         tensor.from_elements %{{.*}} : tensor<1xi32>
-// DET-ALL:         linalg.tensor_collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
+// DET-ALL:         tensor.collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
 // DET-ALL:         linalg.init_tensor [10] : tensor<10xi32>
 // DET-ALL:         linalg.generic {{{.*}}} ins(%{{.*}} : tensor<i32>) outs(%{{.*}} : tensor<10xi32>) {
 // DET-ALL:         ^bb0(%{{.*}}: i32, %{{.*}}: i32):
 // DET-ALL:           linalg.yield %{{.*}} : i32
 // DET-ALL:         } -> tensor<10xi32>
 // DET-ALL:         br ^[[bb1]](%{{.*}} : tensor<10xi32>)
 // DET-ALL:       ^[[bb3]](%{{.*}}: i32)
 // DET-ALL:         tensor.from_elements %{{.*}} : tensor<1xi32>
-// DET-ALL:         linalg.tensor_collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
+// DET-ALL:         tensor.collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
 // DET-ALL:         return %{{.*}} : tensor<i32>
 // DET-ALL:       }
 
 // DET-CF-LABEL: func @main
 // DET-CF-SAME:    (%{{.*}}: tensor<10xi32>, %{{.*}}: tensor<i32>)
 // DET-CF:         br ^[[bb1:.*]](%{{.*}} : tensor<10xi32>)
 // DET-CF:       ^bb1(%{{.*}}: tensor<10xi32>)
 // DET-CF:         %{{.*}} = linalg.generic {{{.*}}} ins(%{{.*}} : tensor<10xi32>) outs(%{{.*}} : tensor<i32>) {

mlir/test/Dialect/Linalg/detensorize_while_pure_cf.mlir

 // RUN: mlir-opt %s -allow-unregistered-dialect -linalg-detensorize | FileCheck %s
 
 #map0 = affine_map<() -> ()>
 
 #attrs = {
   indexing_maps = [#map0, #map0, #map0],
   iterator_types = []
 }
 
 func @main() -> () attributes {} {
   %c0 = arith.constant 0 : i32
   %0 = tensor.from_elements %c0 : tensor<1xi32>
-  %reshaped0 = linalg.tensor_collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
+  %reshaped0 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
   %c10 = arith.constant 10 : i32
   %1 = tensor.from_elements %c10 : tensor<1xi32>
-  %reshaped1 = linalg.tensor_collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
+  %reshaped1 = tensor.collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
   br ^bb1(%reshaped0 : tensor<i32>)
 
 ^bb1(%2: tensor<i32>):  // 2 preds: ^bb0, ^bb2
   %3 = linalg.init_tensor [] : tensor<i1>
   %4 = linalg.generic #attrs
     ins(%2, %reshaped1 : tensor<i32>, tensor<i32>)
     outs(%3 : tensor<i1>) {
     ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):  // no predecessors

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

outs(%shape : tensor<?x1x?x1x?xf32>) {
          linalg.yield %arg3 : f32
        } -> tensor<?x1x?x1x?xf32>
   return %0 : tensor<?x1x?x1x?xf32>
 }
 //   CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
 //   CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> ()>
 //   CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 // CHECK-LABEL: func @drop_one_trip_loops
-//       CHECK: linalg.tensor_collapse_shape %{{.*}} {{\[}}[0, 1], [2]]
+//       CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1], [2]]
 //       CHECK: linalg.generic
 //  CHECK-SAME:   indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP3]]]
 //  CHECK-SAME:   iterator_types = ["parallel", "parallel", "parallel"]
-//       CHECK: linalg.tensor_expand_shape %{{.*}} {{\[}}[0, 1], [2, 3], [4]]
+//       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2, 3], [4]]
 
 // -----
 
 #accesses = [
   affine_map<(i, j, k, l, m) -> (i, k, m)>,
   affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>
 ]
 
%0 = linalg.generic #trait
     outs(%arg0 : tensor<1x1xf32>) {
        ^bb0(%arg1: f32, %arg2: f32) :
          linalg.yield %arg1 : f32
        } -> tensor<1x1xf32>
   return %0 : tensor<1x1xf32>
 }
 //       CHECK: #[[$MAP0:.*]] = affine_map<() -> ()>
 // CHECK-LABEL: func @drop_all_loops
-//       CHECK:   linalg.tensor_collapse_shape %{{.*}} []
+//       CHECK:   tensor.collapse_shape %{{.*}} []
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[$MAP0]], #[[$MAP0]]]
 //  CHECK-SAME:     iterator_types = []
 
 // -----
 
 #map0 = affine_map<(i, j) -> (i, j)>
 #access = [#map0, #map0]
func @leading_dim_1_canonicalization(%arg0: tensor<1x5xf32>, %shape: tensor<5xf32>) -> tensor<5xf32> {
   ^bb0(%arg2: f32, %arg3: f32):     // no predecessors
     linalg.yield %arg2 : f32
   } -> tensor<5xf32>
   return %0 : tensor<5xf32>
 }
 //   CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL: func @leading_dim_1_canonicalization
-//       CHECK:   linalg.tensor_collapse_shape %{{.*}} {{\[}}[0, 1]]
+//       CHECK:   tensor.collapse_shape %{{.*}} {{\[}}[0, 1]]
 //       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 = {
   indexing_maps = #accesses,
   iterator_types = ["parallel", "parallel"],
   library_call = "some_external_fn"
 }
 
 func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>, %shape : tensor<5x5xf32>) -> tensor<5x5xf32>
 {
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1]] : tensor<5xf32> into tensor<1x5xf32>
-  %1 = linalg.tensor_expand_shape %arg1 [[0, 1]] : tensor<5xf32> into tensor<5x1xf32>
+  %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<5xf32> into tensor<1x5xf32>
+  %1 = tensor.expand_shape %arg1 [[0, 1]] : tensor<5xf32> into tensor<5x1xf32>
   %2 = linalg.generic #trait
      ins(%0, %1 : tensor<1x5xf32>, tensor<5x1xf32>)
     outs(%shape : tensor<5x5xf32>) {
        ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
          %3 = arith.addf %arg3, %arg4 : f32
          linalg.yield %3 : f32
        } -> tensor<5x5xf32>
   return %2 : tensor<5x5xf32>
outs(%shape : tensor<?x?xf32>) {
         linalg.yield %arg2 : f32
    } -> 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_collapse_shape %[[ARG0]] []
+//       CHECK:   %[[A:.*]] = tensor.collapse_shape %[[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]]
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 #map1 = affine_map<(d0, d1, d2) -> (d2)>
 func @fold_unit_dim_tensor_reshape_op(%arg0 : tensor<5xf32>) -> tensor<2x5xf32>
 {
   %1 = linalg.init_tensor [1, 2, 5] : tensor<1x2x5xf32>
   %2 = linalg.generic {i64, indexing_maps = [#map1, #map0],
     iterator_types = ["parallel", "parallel", "parallel"]}
     ins(%arg0 : tensor<5xf32>) outs(%1 : tensor<1x2x5xf32>) {
     ^bb0(%arg1: f32, %arg2: f32):  // no predecessors
       linalg.yield %arg1 : f32
     } -> tensor<1x2x5xf32>
-  %3 = linalg.tensor_collapse_shape %2 [[0, 1], [2]]
+  %3 = tensor.collapse_shape %2 [[0, 1], [2]]
     : tensor<1x2x5xf32> into tensor<2x5xf32>
   return %3 : tensor<2x5xf32>
 }
 // CHECK-LABEL: func @fold_unit_dim_tensor_reshape_op
 //       CHECK:   %[[RESULT:.+]] = linalg.generic
 //       CHECK:   return %[[RESULT]]
 
 // -----

 
 
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0)>
 //   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> ()>
 
 //       CHECK: func @fold_unit_dim_for_init_tensor
 
 
-//       CHECK: %[[INPUT_RESHAPE:.+]] = linalg.tensor_collapse_shape %{{.+}} {{\[}}[0, 1]] : tensor<1x1000xf32> into tensor<1000xf32>
+//       CHECK: %[[INPUT_RESHAPE:.+]] = tensor.collapse_shape %{{.+}} {{\[}}[0, 1]] : tensor<1x1000xf32> into tensor<1000xf32>
 //       CHECK: %[[INIT:.+]] = linalg.init_tensor [] : tensor<f32>
 //       CHECK: %[[FILL:.+]] = linalg.fill(%cst, %[[INIT]]) : f32, tensor<f32> -> tensor<f32>
 //       CHECK: %[[GENERIC:.+]] = linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP1]], #[[MAP2]]]
 //  CHECK-SAME:     iterator_types = ["reduction"]
 //  CHECK-SAME:   ins(%[[INPUT_RESHAPE]] : tensor<1000xf32>)
 //  CHECK-SAME:   outs(%[[FILL]] : tensor<f32>)
-//       CHECK: %[[GENERIC_RESHAPE:.+]] = linalg.tensor_expand_shape %[[GENERIC]] [] : tensor<f32> into tensor<1xf32>
+//       CHECK: %[[GENERIC_RESHAPE:.+]] = tensor.expand_shape %[[GENERIC]] [] : tensor<f32> into tensor<1xf32>
 //       CHECK: return %[[GENERIC_RESHAPE:.+]] : tensor<1xf32>
 
 
 // -----
 
 func @fold_slice(
     %arg0 : tensor<1x?x?x1x?x1x1xf32>, %arg1 : tensor<1x?x?x?x?x1x1xf32>,
     %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : index,
     %arg6 : index, %arg7 : index) -> (tensor<1x?x?x1x?x1x1xf32>, tensor<1x?x?x1x?x1x1xf32>) {
   %0 = tensor.extract_slice %arg0[0, %arg2, %arg3, 0, %arg4, 0, 0]
                              [1, %arg5, %arg6, 1, %arg7, 1, 1] [1, 1, 1, 1, 1, 1, 1] :
       tensor<1x?x?x1x?x1x1xf32> to tensor<1x?x?x1x?x1x1xf32>
   %1 = tensor.extract_slice %arg1[%arg2, 0, %arg3, 0, 0, %arg4, 0]
                              [1, %arg5, %arg6, 1, %arg7, 1, 1] [1, 1, 1, 1, 1, 1, 1] :
       tensor<1x?x?x?x?x1x1xf32> to tensor<1x?x?x1x?x1x1xf32>
   return %0, %1 : tensor<1x?x?x1x?x1x1xf32>, tensor<1x?x?x1x?x1x1xf32>
 }
 //      CHECK: func @fold_slice
 // CHECK-SAME:   %[[ARG0:.+]]: tensor<1x?x?x1x?x1x1xf32>
 // CHECK-SAME:   %[[ARG1:.+]]: tensor<1x?x?x?x?x1x1xf32>
 //      CHECK:   %[[SLICE1:.+]] = tensor.extract_slice %[[ARG0]]
 // CHECK-SAME:       to tensor<?x?x?xf32>
-//      CHECK:   %[[RESULT1:.+]] = linalg.tensor_expand_shape %[[SLICE1]]
+//      CHECK:   %[[RESULT1:.+]] = tensor.expand_shape %[[SLICE1]]
 // CHECK-SAME:       [0, 1], [2], [3, 4, 5, 6]
 //      CHECK:   %[[SLICE2:.+]] = tensor.extract_slice %[[ARG1]]
 // CHECK-SAME:       to tensor<?x?x?xf32>
-//      CHECK:   %[[RESULT2:.+]] = linalg.tensor_expand_shape %[[SLICE2]]
+//      CHECK:   %[[RESULT2:.+]] = tensor.expand_shape %[[SLICE2]]
 // CHECK-SAME:       [0, 1], [2], [3, 4, 5, 6]
 //      CHECK:   return %[[RESULT1]], %[[RESULT2]]
 
 // -----
 
 func @unit_dim_for_reduction(%arg0: tensor<1x?x1x?xf32>) -> tensor<1x?xf32> {
   %cst = arith.constant 1.000000e+00 : f32
   %c3 = arith.constant 3 : index
^bb0(%arg1: f32, %arg2: f32):  // no predecessors
     linalg.yield %4 : f32
   } -> tensor<1x?xf32>
   return %3 : tensor<1x?xf32>
 }
 //  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)>
 //  CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d0)>
 //      CHECK: func @unit_dim_for_reduction
 // CHECK-SAME:   %[[ARG0:.+]]: tensor<1x?x1x?xf32>
-//  CHECK-DAG:   %[[RESHAPE:.+]] = linalg.tensor_collapse_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]]
+//  CHECK-DAG:   %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]]
 //      CHECK:   %[[INIT:.+]] = linalg.init_tensor [%{{.+}}] : tensor<?xf32>
 //      CHECK:   %[[FILL:.+]] = linalg.fill(%{{.+}}, %[[INIT]])
 //      CHECK:   %[[RESULT:.+]] = linalg.generic
 // CHECK-SAME:     indexing_maps = [#[[MAP2]], #[[MAP3]]]
 // CHECK-SAME:     iterator_types = ["parallel", "reduction"]
 // CHECK-SAME:     ins(%[[RESHAPE]] : tensor<?x?xf32>)
 // CHECK-SAME:     outs(%[[FILL]] : tensor<?xf32>)
-//      CHECK:   %[[RESULT_RESHAPE:.+]] = linalg.tensor_expand_shape %[[RESULT]] {{\[}}[0, 1]]
+//      CHECK:   %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]]
 //      CHECK:   return %[[RESULT_RESHAPE]]
 
 // -----
 
 func @unit_dim_for_both_reduction(%arg0: tensor<1x?x1x1xf32>) -> tensor<1x1xf32> {
   %cst = arith.constant 1.000000e+00 : f32
   %c3 = arith.constant 3 : index
   %1 = linalg.init_tensor [1, 1] : tensor<1x1xf32>
   %2 = linalg.fill(%cst, %1) : f32, tensor<1x1xf32> -> tensor<1x1xf32>
   %3 = linalg.generic {
     indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,
                      affine_map<(d0, d1, d2, d3) -> (d0, d1)>],
     iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
     ins(%arg0 : tensor<1x?x1x1xf32>)
     outs(%2 : tensor<1x1xf32>) {
   ^bb0(%arg1: f32, %arg2: f32):  // no predecessors
     %4 = arith.addf %arg1, %arg2 : f32
     linalg.yield %4 : f32
   } -> tensor<1x1xf32>
   return %3 : tensor<1x1xf32>
 }
 //  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> (d0)>
 //      CHECK: func @unit_dim_for_both_reduction
 // CHECK-SAME:   %[[ARG0:.+]]: tensor<1x?x1x1xf32>
-//  CHECK-DAG:   %[[RESHAPE:.+]] = linalg.tensor_collapse_shape %[[ARG0]] {{\[}}[0, 1, 2, 3]
+//  CHECK-DAG:   %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2, 3]
 //      CHECK:   %[[INIT:.+]] = linalg.init_tensor [1] : tensor<1xf32>
 //      CHECK:   %[[FILL:.+]] = linalg.fill(%{{.+}}, %[[INIT]])
 //      CHECK:   %[[RESULT:.+]] = linalg.generic
 // CHECK-SAME:     indexing_maps = [#[[MAP2]], #[[MAP2]]]
 // CHECK-SAME:     iterator_types = ["parallel"]
 // CHECK-SAME:     ins(%[[RESHAPE]] : tensor<?xf32>)
 // CHECK-SAME:     outs(%[[FILL]] : tensor<1xf32>)
-//      CHECK:   %[[RESULT_RESHAPE:.+]] = linalg.tensor_expand_shape %[[RESULT]] {{\[}}[0, 1]]
+//      CHECK:   %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]]
 //      CHECK:   return %[[RESULT_RESHAPE]]
 
 // -----
 
 func @unit_dim_for_reduction_inner(%arg0: tensor<?x1x?x1xf32>) -> tensor<?x1xf32> {
   %cst = arith.constant 1.000000e+00 : f32
   %c2 = arith.constant 2 : index
   %0 = tensor.dim %arg0, %c2 : tensor<?x1x?x1xf32>
^bb0(%arg1: f32, %arg2: f32):  // no predecessors
     linalg.yield %4 : f32
   } -> tensor<?x1xf32>
   return %3 : tensor<?x1xf32>
 }
 //  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)>
 //  CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d0)>
 //      CHECK: func @unit_dim_for_reduction_inner
 // CHECK-SAME:   %[[ARG0:.+]]: tensor<?x1x?x1xf32>
-//  CHECK-DAG:   %[[RESHAPE:.+]] = linalg.tensor_collapse_shape %[[ARG0]] {{\[}}[0, 1], [2, 3]]
+//  CHECK-DAG:   %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1], [2, 3]]
 //      CHECK:   %[[INIT:.+]] = linalg.init_tensor [%{{.+}}] : tensor<?xf32>
 //      CHECK:   %[[FILL:.+]] = linalg.fill(%{{.+}}, %[[INIT]])
 //      CHECK:   %[[RESULT:.+]] = linalg.generic
 // CHECK-SAME:     indexing_maps = [#[[MAP2]], #[[MAP3]]]
 // CHECK-SAME:     iterator_types = ["parallel", "reduction"]
 // CHECK-SAME:     ins(%[[RESHAPE]] : tensor<?x?xf32>)
 // CHECK-SAME:     outs(%[[FILL]] : tensor<?xf32>)
-//      CHECK:   %[[RESULT_RESHAPE:.+]] = linalg.tensor_expand_shape %[[RESULT]] {{\[}}[0, 1]]
+//      CHECK:   %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]]
 //      CHECK:   return %[[RESULT_RESHAPE]]
 
 // -----
 
 func @slice_unit_dims(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
   %0 = tensor.extract_slice %arg0[0, 2] [1, 1] [1, 1] : tensor<1x3xf32> to tensor<1x1xf32>
   return %0 : tensor<1x1xf32>
 }
 // CHECK-LABEL: func @slice_unit_dims
 //       CHECK:   %[[SLICE:.+]] = tensor.extract_slice
 //  CHECK-SAME:     tensor<1x3xf32> to tensor<f32>
-//       CHECK:   %[[RESULT:.+]] = linalg.tensor_expand_shape %[[SLICE]] []
+//       CHECK:   %[[RESULT:.+]] = tensor.expand_shape %[[SLICE]] []
 //       CHECK:   return %[[RESULT]]
 
 // -----
 
 func @insert_slice_unit_dims(%arg0: tensor<1x3xf32>, %arg1: tensor<1x1xf32>) -> tensor<1x3xf32> {
   %0 = tensor.insert_slice %arg1 into %arg0[0, 2] [1, 1] [1, 1] : tensor<1x1xf32> into tensor<1x3xf32>
   return %0 : tensor<1x3xf32>
 }
 // CHECK-LABEL: func @insert_slice_unit_dims
-//       CHECK:   %[[RESHAPE:.+]] = linalg.tensor_collapse_shape %{{.+}} []
+//       CHECK:   %[[RESHAPE:.+]] = tensor.collapse_shape %{{.+}} []
 //       CHECK:   %[[RESULT:.+]] = tensor.insert_slice %[[RESHAPE]]
 //  CHECK-SAME:     tensor<f32> into tensor<1x3xf32>
 //       CHECK:   return %[[RESULT]]
 
 // -----
 
 #accesses = [
   affine_map<(i, j, k, l, m) -> (i, k, m)>,

mlir/test/Dialect/Linalg/fusion-push-reshape.mlir

 // RUN: mlir-opt %s -test-linalg-push-reshape -split-input-file | FileCheck %s
 
 // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1) -> (d1)>
 
 // CHECK-LABEL: func @reshape
 // CHECK-SAME: (%[[A:.*]]: tensor<?x16xf32>, %[[B:.*]]: tensor<16xf32>, %[[INIT:.*]]: tensor<?x112x16xf32>)
-//      CHECK: %[[RI:.*]] = linalg.tensor_collapse_shape %[[INIT]] {{\[}}[0, 1], [2]] : tensor<?x112x16xf32> into tensor<?x16xf32>
+//      CHECK: %[[RI:.*]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1], [2]] : tensor<?x112x16xf32> into tensor<?x16xf32>
 //      CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP3]], #[[$MAP2]]],
 // CHECK-SAME: iterator_types = ["parallel", "parallel"]}
 // CHECK-SAME: ins(%[[A]], %[[B]] : tensor<?x16xf32>, tensor<16xf32>) outs(%[[RI]] : tensor<?x16xf32>)
-//      CHECK: %[[RR:.*]] = linalg.tensor_expand_shape %[[R]] {{\[}}[0, 1], [2]] : tensor<?x16xf32> into tensor<?x112x16xf32>
+//      CHECK: %[[RR:.*]] = tensor.expand_shape %[[R]] {{\[}}[0, 1], [2]] : tensor<?x16xf32> into tensor<?x112x16xf32>
 //      CHECK: return %[[RR]] : tensor<?x112x16xf32>
 func @reshape(%A: tensor<?x16xf32>, %B: tensor<16xf32>, %init: tensor<?x112x16xf32>) -> tensor<?x112x16xf32> {
-  %0 = linalg.tensor_expand_shape %A [[0, 1], [2]]
+  %0 = tensor.expand_shape %A [[0, 1], [2]]
       : tensor<?x16xf32> into tensor<?x112x16xf32>
   %2 = linalg.generic {indexing_maps = [
     affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d2)>,
     affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
     iterator_types = ["parallel", "parallel", "parallel"]}
   ins(%0, %B : tensor<?x112x16xf32>, tensor<16xf32>)
   outs(%init : tensor<?x112x16xf32>) {
   ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):  // no predecessors
     %s = arith.subf %arg1, %arg2 : f32
     linalg.yield %s : f32
   } -> tensor<?x112x16xf32>
   return %2 : tensor<?x112x16xf32>
 }
 
 // -----
 
 // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1) -> (d1)>
 
 // CHECK-LABEL: func @reshape_multiple
 // CHECK-SAME: (%[[A:.*]]: tensor<12544x16xf32>, %[[B:.*]]: tensor<12544x16xf32>, %[[C:.*]]: tensor<16xf32>)
 //      CHECK: %[[I:.*]] = linalg.init_tensor [112, 112, 16] : tensor<112x112x16xf32>
-//      CHECK: %[[RI:.*]] = linalg.tensor_collapse_shape %[[I]] {{\[}}[0, 1], [2]] : tensor<112x112x16xf32> into tensor<12544x16xf32>
+//      CHECK: %[[RI:.*]] = tensor.collapse_shape %[[I]] {{\[}}[0, 1], [2]] : tensor<112x112x16xf32> into tensor<12544x16xf32>
 //      CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP2]], #[[$MAP3]], #[[$MAP2]]],
 // CHECK-SAME: iterator_types = ["parallel", "parallel"]}
 // CHECK-SAME: ins(%[[A]], %[[B]], %[[C]] : tensor<12544x16xf32>, tensor<12544x16xf32>, tensor<16xf32>) outs(%[[RI]] : tensor<12544x16xf32>)
-//      CHECK: %[[RR:.*]] = linalg.tensor_expand_shape %[[R]] {{\[}}[0, 1], [2]] : tensor<12544x16xf32> into tensor<112x112x16xf32>
+//      CHECK: %[[RR:.*]] = tensor.expand_shape %[[R]] {{\[}}[0, 1], [2]] : tensor<12544x16xf32> into tensor<112x112x16xf32>
 //      CHECK: return %[[RR]] : tensor<112x112x16xf32>
 func @reshape_multiple(%A: tensor<12544x16xf32>, %B: tensor<12544x16xf32>,
   %C: tensor<16xf32>) -> tensor<112x112x16xf32> {
-  %0 = linalg.tensor_expand_shape %A [[0, 1], [2]]
+  %0 = tensor.expand_shape %A [[0, 1], [2]]
       : tensor<12544x16xf32> into tensor<112x112x16xf32>
-  %1 = linalg.tensor_expand_shape %B [[0, 1], [2]]
+  %1 = tensor.expand_shape %B [[0, 1], [2]]
       : tensor<12544x16xf32> into tensor<112x112x16xf32>
   %2 = linalg.init_tensor [112, 112, 16] : tensor<112x112x16xf32>
   %3 = linalg.generic {indexing_maps = [
     affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
     affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
     affine_map<(d0, d1, d2) -> (d2)>,
     affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
     iterator_types = ["parallel", "parallel", "parallel"]}
   ins(%0, %1, %C : tensor<112x112x16xf32>, tensor<112x112x16xf32>, tensor<16xf32>)
   outs(%2 : tensor<112x112x16xf32>) {
   ^bb0(%arg1: f32, %arg2: f32, %arg3: f32, %arg4: f32):  // no predecessors
     %s = arith.subf %arg1, %arg2 : f32
     %m = arith.mulf %s, %arg3 : f32
     linalg.yield %m : f32
   } -> tensor<112x112x16xf32>
   return %3 : tensor<112x112x16xf32>
 }
 
 // -----
 
 // Negative test, since the second source is broadcasted from d1 we cannot merge
 // d0 and d1 dimensions
 // CHECK-LABEL: func @reshape_negative
-// CHECK: linalg.tensor_expand_shape {{.*}} : tensor<12544x16xf32> into tensor<112x112x16xf32>
+// CHECK: tensor.expand_shape {{.*}} : tensor<12544x16xf32> into tensor<112x112x16xf32>
 // CHECK: linalg.generic
 // CHECK: } -> tensor<112x112x16xf32>
 func @reshape_negative(%A: tensor<12544x16xf32>, %B: tensor<112xf32>) -> tensor<112x112x16xf32> {
-  %20 = linalg.tensor_expand_shape %A [[0, 1], [2]]
+  %20 = tensor.expand_shape %A [[0, 1], [2]]
       : tensor<12544x16xf32> into tensor<112x112x16xf32>
   %21 = linalg.init_tensor [112, 112, 16] : tensor<112x112x16xf32>
   %22 = linalg.generic {indexing_maps = [
     affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d1)>,
     affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
     iterator_types = ["parallel", "parallel", "parallel"]}
   ins(%20, %B : tensor<112x112x16xf32>, tensor<112xf32>)
   outs(%21 : tensor<112x112x16xf32>) {
   ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):  // no predecessors
     %s = arith.subf %arg1, %arg2 : f32
     linalg.yield %s : f32
   } -> tensor<112x112x16xf32>
   return %22 : tensor<112x112x16xf32>
 }
 
 // -----
 
 func @type_correctness(%arg0 : tensor<6x5xi32>, %arg1 : tensor<5xf32>,
     %arg2 : tensor<5xf32>) -> tensor<2x3x5xf32> {
   %cst_6 = arith.constant 1.000000e+00 : f32
   %cst_7 = arith.constant 7.000000e+00 : f32
   %cst_8 = arith.constant 1.1920929E-7 : f32
-  %25 = linalg.tensor_expand_shape %arg0 [[0, 1], [2]]
+  %25 = tensor.expand_shape %arg0 [[0, 1], [2]]
       : tensor<6x5xi32> into tensor<2x3x5xi32>
   %26 = linalg.init_tensor [2, 3, 5] : tensor<2x3x5xf32>
   %28 = linalg.generic {
       indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                        affine_map<(d0, d1, d2) -> (d2)>,
                        affine_map<(d0, d1, d2) -> (d2)>,
                        affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
       iterator_types = ["parallel", "parallel", "parallel"]}
%28 = linalg.generic {
         linalg.yield %34 : f32
       } -> tensor<2x3x5xf32>
   return %28 : tensor<2x3x5xf32>
 }
 // CHECK-LABEL: func @type_correctness
 //       CHECK:   %[[OP:.+]] = linalg.generic
 //  CHECK-SAME:   ins(%{{.+}}, %{{.+}}, %{{.+}} : tensor<6x5xi32>, tensor<5xf32>, tensor<5xf32>)
 //  CHECK-SAME:   outs(%{{.+}} : tensor<6x5xf32>)
-//       CHECK:   linalg.tensor_expand_shape %[[OP]]
+//       CHECK:   tensor.expand_shape %[[OP]]
 //  CHECK-SAME:   tensor<6x5xf32> into tensor<2x3x5xf32>

mlir/test/Dialect/Linalg/invalid.mlir

 {
   // expected-error @+1 {{expected 2 dynamic sizes values}}
   %1 = "linalg.init_tensor"(%arg0) {static_sizes = [4, -1, -1, 5]} : (index) -> tensor<4x?x?x5xf32>
   return
 }
 
 // -----
 
-func @illegal_expanding_reshape_dynamic_tensor
-  (%arg0: tensor<?x?x?xf32>) -> tensor<?x?x?x4x?xf32>
-{
-  // expected-error @+1 {{invalid to have a single dimension (2) expanded into multiple dynamic dims (2,4)}}
-  %0 = linalg.tensor_expand_shape %arg0 [[0], [1], [2, 3, 4]]
-      : tensor<?x?x?xf32> into tensor<?x?x?x4x?xf32>
-  return %0 : tensor<?x?x?x4x?xf32>
-}
-
-// -----
-
-
-func @illegal_expanding_reshape_static_tensor
-  (%arg0: tensor<2x3x20xf32>) -> tensor<2x3x2x4x5xf32>
-{
-  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
-  %0 = linalg.tensor_expand_shape %arg0 [[0], [1], [2, 3, 4]]
-      : tensor<2x3x20xf32> into tensor<2x3x2x4x5xf32>
-  return %0 : tensor<2x3x2x4x5xf32>
-}
-
-// -----
-
-func @illegal_collapsing_reshape_static_tensor
-  (%arg0: tensor<2x3x2x4x5xf32>) -> tensor<2x3x20xf32>
-{
-  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
-  %0 = linalg.tensor_collapse_shape %arg0 [[0], [1], [2, 3, 4]]
-      : tensor<2x3x2x4x5xf32> into tensor<2x3x20xf32>
-  return %0 : tensor<2x3x20xf32>
-}
-
-// -----
-
-func @illegal_expanding_reshape_mixed_tensor(%arg0 : tensor<?x?xf32>) -> tensor<?x4x5xf32>
-{
-  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1], [2]]
-      : tensor<?x?xf32> into tensor<?x4x5xf32>
-  return %0 : tensor<?x4x5xf32>
-}
-
-// -----
-
-func @illegal_expanding_reshape_mixed_tensor_2(%arg0 : tensor<?x?xf32>) -> tensor<?x4x5xf32>
-{
-  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
-  %0 = linalg.tensor_expand_shape %arg0 [[0], [1, 2]]
-      : tensor<?x?xf32> into tensor<?x4x5xf32>
-  return %0 : tensor<?x4x5xf32>
-}
-
-// -----
-
-func @illegal_collapsing_reshape_mixed_tensor(%arg0 : tensor<?x4x5xf32>) -> tensor<?x?xf32>
-{
-  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
-  %0 = linalg.tensor_collapse_shape %arg0 [[0, 1], [2]]
-      : tensor<?x4x5xf32> into tensor<?x?xf32>
-  return %0 : tensor<?x?xf32>
-}
-
-// -----
-
-func @illegal_collapsing_reshape_mixed_tensor_2(%arg0 : tensor<?x4x5xf32>) -> tensor<?x?xf32>
-{
-  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
-  %0 = linalg.tensor_collapse_shape %arg0 [[0], [1, 2]]
-      : tensor<?x4x5xf32> into tensor<?x?xf32>
-  return %0 : tensor<?x?xf32>
-}
-
-// -----
 
 func @pad_result_type(%arg0: tensor<?x2x3x4xi32>, %arg1: index, %arg2: i32) -> tensor<?x?x?x8xf32> {
   // expected-error @+1 {{specified type 'tensor<?x?x?x8xf32>' does not match the inferred type 'tensor<?x?x?x9xi32>}}
   %0 = linalg.pad_tensor %arg0 low[1, %arg1, 2, 2] high[1, 2, %arg1, 3] {
   ^bb0(%arg3: index, %arg4: index):  // no predecessors
     linalg.yield %arg2 : i32
   } : tensor<?x2x3x4xi32> to tensor<?x?x?x8xf32>
   return %0 : tensor<?x?x?x8xf32>

mlir/test/Dialect/Linalg/reshape_control_fusion.mlir

 // RUN: mlir-opt -test-linalg-control-fusion-by-expansion %s -split-input-file | FileCheck %s
 
 func @control_producer_reshape_fusion(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?xf32>) -> tensor<?x?xf32> {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
-  %0 = linalg.tensor_collapse_shape %arg0 [[0, 1], [2]] : tensor<?x?x?xf32> into tensor<?x?xf32>
+  %0 = tensor.collapse_shape %arg0 [[0, 1], [2]] : tensor<?x?x?xf32> into tensor<?x?xf32>
   %d0 = tensor.dim %0, %c0 : tensor<?x?xf32>
   %d1 = tensor.dim %0, %c1 : tensor<?x?xf32>
   %init = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32>
   %1 = linalg.generic {
       indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0, d1)>],
       iterator_types = ["parallel", "parallel"]}
       ins(%0, %arg1 : tensor<?x?xf32>, tensor<?xf32>)
       outs(%init : tensor<?x?xf32>) {
       ^bb0(%arg2 : f32, %arg3:f32, %arg4 : f32):
         %2 = arith.addf %arg2, %arg3 : f32
         linalg.yield %2 : f32
       } -> tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
 //  CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
 //  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d1)>
 //      CHECK: func @control_producer_reshape_fusion
 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
 // CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?xf32>
 //  CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
 //  CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
-//      CHECK:   %[[RESHAPE:.+]] = linalg.tensor_collapse_shape %[[ARG0]]
+//      CHECK:   %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]]
 // CHECK-SAME:       {{\[}}[0, 1], [2]{{\]}} : tensor<?x?x?xf32> into tensor<?x?xf32>
 //      CHECK:   %[[RESULT:.+]] = linalg.generic
 // CHECK-SAME:       indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP0]]]
 // CHECK-SAME:       ins(%[[RESHAPE]], %[[ARG1]] : tensor<?x?xf32>, tensor<?xf32>)
 //      CHECK:   return %[[RESULT]]
 
 // -----
 
 func @control_consumer_reshape_fusion(%arg0 : tensor<1x?x?xf32>, %arg1 : tensor<1x?x?xf32>) -> tensor<1x?x?xf32> {
   %c1 = arith.constant 1 : index
   %c2 = arith.constant 2 : index
   %cst = arith.constant 0.0 : f32
   %d0 = tensor.dim %arg0, %c1 : tensor<1x?x?xf32>
   %d1 = tensor.dim %arg1, %c2 : tensor<1x?x?xf32>
   %init = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32>
   %fill = linalg.generic {
       indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
       iterator_types = ["parallel", "parallel"]}
       outs(%init : tensor<?x?xf32>) {
       ^bb0(%arg2: f32):
         linalg.yield %cst : f32
       } -> tensor<?x?xf32>
-  %0 = linalg.tensor_expand_shape %fill [[0, 1], [2]] : tensor<?x?xf32> into tensor<1x?x?xf32>
+  %0 = tensor.expand_shape %fill [[0, 1], [2]] : tensor<?x?xf32> into tensor<1x?x?xf32>
   %1 = linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x?x?xf32>, tensor<1x?x?xf32>)
       outs(%0 : tensor<1x?x?xf32>) -> tensor<1x?x?xf32>
   return %1 : tensor<1x?x?xf32>
 }
 //  CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)
 //      CHECK: func @control_consumer_reshape_fusion
 //      CHECK:   %[[FILL:.+]] = linalg.generic
 // CHECK-SAME:       indexing_maps = [#[[MAP]]]
 // CHECK-SAME:       outs(%{{.+}} : tensor<1x?x?xf32>)
 //      CHECK:   linalg.batch_matmul
 // CHECK-SAME:       outs(%[[FILL]] : tensor<1x?x?xf32>)

mlir/test/Dialect/Linalg/reshape_fusion.mlir

 // RUN: mlir-opt %s -linalg-fuse-elementwise-ops="allow-folding-unit-dim-reshapes=false" -split-input-file | FileCheck %s
 // RUN: mlir-opt %s -linalg-fuse-elementwise-ops="allow-folding-unit-dim-reshapes=true" -split-input-file | FileCheck %s --check-prefix=FOLDUNITDIM
 #map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
 #map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
 #map2 = affine_map<(d0, d1, d2) -> ()>
 func @generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x4x?xf32>,
                                          %arg1 : tensor<?x?x?xf32>,
                                          %arg2 : f32) ->
                                          tensor<?x?x?xf32>
 {
-  %0 = linalg.tensor_collapse_shape %arg0 [[0], [1, 2], [3]] :
+  %0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3]] :
     tensor<?x?x4x?xf32> into tensor<?x?x?xf32>
   %1 = linalg.generic {
      indexing_maps = [#map0, #map1, #map2, #map1],
      iterator_types = ["parallel", "parallel", "parallel"]}
        ins(%0, %arg1, %arg2 : tensor<?x?x?xf32>, tensor<?x?x?xf32>, f32)
        outs(%0 : tensor<?x?x?xf32>) {
     ^bb0(%arg3: f32, %arg4: f32, %arg5: f32, %s: f32):       // no predecessors
       %1 = arith.mulf %arg3, %arg4 : f32
       %2 = arith.addf %1, %arg5 : f32
       linalg.yield %2 : f32
   } -> tensor<?x?x?xf32>
   return %1 : tensor<?x?x?xf32>
 }
 
 //  CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)>
 //  CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d0, d1)>
 //  CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1, d2, d3) -> ()>
 //      CHECK: func @generic_op_reshape_producer_fusion
 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x4x?xf32>
 // CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
 // CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: f32
-//      CHECK:   %[[T0:.+]] = linalg.tensor_collapse_shape %[[ARG0]]
+//      CHECK:   %[[T0:.+]] = tensor.collapse_shape %[[ARG0]]
 // CHECK-SAME:     [0], [1, 2], [3]
-//      CHECK:   %[[T1:.+]] = linalg.tensor_expand_shape %[[ARG1]]
+//      CHECK:   %[[T1:.+]] = tensor.expand_shape %[[ARG1]]
 // CHECK-SAME:     [0], [1], [2, 3]
 //      CHECK:   %[[T3:.+]] = linalg.generic
 // CHECK-SAME:     indexing_maps = [#[[MAP5]], #[[MAP6]], #[[MAP7]], #[[MAP6]]]
 // CHECK-SAME:     ["parallel", "parallel", "parallel", "parallel"]
 // CHECK-SAME:     ins(%[[ARG0]], %[[T1]], %[[ARG2]] : tensor<?x?x4x?xf32>, tensor<?x?x?x4xf32>, f32)
 // CHECK-SAME:     outs(%{{.+}} : tensor<?x?x?x4xf32>)
-//      CHECK:   %[[T4:.+]] = linalg.tensor_collapse_shape %[[T3]]
+//      CHECK:   %[[T4:.+]] = tensor.collapse_shape %[[T3]]
 // CHECK-SAME:     [0], [1], [2, 3]
 // CHECK-SAME:     tensor<?x?x?x4xf32> into tensor<?x?x?xf32>
 //      CHECK:   return %[[T4]]
 
 // -----
 
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
 #map1 = affine_map<(d0, d1) -> ()>
 func @generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?xf32>,
                                          %arg1 : tensor<?x?xf32>,
                                          %arg2 : f32) ->
                                          tensor<?x4x?x5xf32>
 {
   %0 = linalg.generic {
      indexing_maps = [#map0, #map0, #map1, #map0],
      iterator_types = ["parallel", "parallel"]}
        ins(%arg0, %arg1, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>, f32)
        outs(%arg0 : tensor<?x?xf32>) {
     ^bb0(%arg3: f32, %arg4: f32, %arg5: f32, %s: f32):       // no predecessors
       %1 = arith.mulf %arg3, %arg4 : f32
       %2 = arith.addf %1, %arg5 : f32
       linalg.yield %2 : f32
   } -> tensor<?x?xf32>
-  %1 = linalg.tensor_expand_shape %0 [[0], [1, 2, 3]] :
+  %1 = tensor.expand_shape %0 [[0], [1, 2, 3]] :
     tensor<?x?xf32> into tensor<?x4x?x5xf32>
   return %1 : tensor<?x4x?x5xf32>
 }
 
 //  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 //  CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> ()>
 
 //      CHECK: func @generic_op_reshape_consumer_fusion
 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
 // CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
 // CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: f32
-//      CHECK:   %[[T0:.+]] = linalg.tensor_expand_shape %[[ARG0]]
+//      CHECK:   %[[T0:.+]] = tensor.expand_shape %[[ARG0]]
 // CHECK-SAME:     [0], [1, 2, 3]
 // CHECK-SAME:     tensor<?x?xf32> into tensor<?x4x?x5xf32>
-//      CHECK:   %[[T1:.+]] = linalg.tensor_expand_shape %[[ARG1]]
+//      CHECK:   %[[T1:.+]] = tensor.expand_shape %[[ARG1]]
 // CHECK-SAME:     [0], [1, 2, 3]
 // CHECK-SAME:     tensor<?x?xf32> into tensor<?x4x?x5xf32>
 //      CHECK:   %[[T3:.+]] = linalg.generic
 // CHECK-SAME:     indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP3]], #[[MAP2]]]
 // CHECK-SAME:     ["parallel", "parallel", "parallel", "parallel"]
 // CHECK-SAME:     ins(%[[T0]], %[[T1]], %[[ARG2]] : tensor<?x4x?x5xf32>, tensor<?x4x?x5xf32>, f32)
 // CHECK-SAME:     outs(%{{.+}} : tensor<?x4x?x5xf32>)
 //      CHECK:   return %[[T3]] : tensor<?x4x?x5xf32>
%c = linalg.generic {
                           affine_map<(d0, d1, d2) -> (d0, d2, d1)>],
          iterator_types = ["parallel", "parallel", "parallel"]}
           ins(%a, %b : tensor<?x?x?xf32>, tensor<?x?xf32>)
          outs(%a : tensor<?x?x?xf32>) {
        ^bb0(%arg0 : f32, %arg1: f32, %s: f32):
          %1 = arith.addf %arg0, %arg1 : f32
          linalg.yield %1 : f32
        } -> tensor<?x?x?xf32>
-  %d = linalg.tensor_expand_shape %c [[0, 1], [2], [3, 4, 5]]
+  %d = tensor.expand_shape %c [[0, 1], [2], [3, 4, 5]]
        : tensor<?x?x?xf32> into tensor<?x2x?x3x4x?xf32>
   return %d : tensor<?x2x?x3x4x?xf32>
 }
 //  CHECK-DAG: #[[MAP8:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)>
 //  CHECK-DAG: #[[MAP9:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)>
 //  CHECK-DAG: #[[MAP10:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)>
 //      CHECK: func @reshape_as_consumer_permutation
 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
 // CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
-//      CHECK:   %[[T0:.+]] = linalg.tensor_expand_shape %[[ARG0]]
+//      CHECK:   %[[T0:.+]] = tensor.expand_shape %[[ARG0]]
 // CHECK-SAME:     [0, 1, 2], [3, 4], [5]
 // CHECK-SAME:     tensor<?x?x?xf32> into tensor<3x4x?x?x2x?xf32>
-//      CHECK:   %[[T1:.+]] = linalg.tensor_expand_shape %[[ARG1]]
+//      CHECK:   %[[T1:.+]] = tensor.expand_shape %[[ARG1]]
 // CHECK-SAME:     [0, 1, 2], [3]
 // CHECK-SAME:     tensor<?x?xf32> into tensor<3x4x?x?xf32>
 //      CHECK:   %[[T3:.+]] = linalg.generic
 // CHECK-SAME:     indexing_maps = [#[[MAP8]], #[[MAP9]], #[[MAP10]]]
 // CHECK-SAME:     ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]
 // CHECK-SAME:     ins(%[[T0]], %[[T1]] : tensor<3x4x?x?x2x?xf32>, tensor<3x4x?x?xf32>)
 // CHECK-SAME:     outs(%{{.+}} : tensor<?x2x?x3x4x?xf32>)
 //      CHECK:   return %[[T3]] : tensor<?x2x?x3x4x?xf32>
%1 = linalg.generic {
      indexing_maps = [#map0, #map0, #map0],
      iterator_types = ["parallel", "parallel"]}
        ins(%arg0, %cst : tensor<264x4xf32>, tensor<264x4xf32>)
        outs(%0 : tensor<264x4xf32>) {
     ^bb0(%arg1: f32, %arg2: f32, %s: f32):  // no predecessors
       %2 = arith.mulf %arg1, %arg2 : f32
       linalg.yield %2 : f32
     } -> tensor<264x4xf32>
-  %2 = linalg.tensor_expand_shape %1 [[0, 1], [2]] :
+  %2 = tensor.expand_shape %1 [[0, 1], [2]] :
     tensor<264x4xf32> into tensor<8x33x4xf32>
   return %2 : tensor<8x33x4xf32>
 }
 
 //  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 //      CHECK: func @generic_op_reshape_consumer_static
 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<264x4xf32>
-//      CHECK:   %[[T0:.+]] = linalg.tensor_expand_shape %[[ARG0]]
+//      CHECK:   %[[T0:.+]] = tensor.expand_shape %[[ARG0]]
 // CHECK-SAME:     [0, 1], [2]
 // CHECK-SAME:     tensor<264x4xf32> into tensor<8x33x4xf32>
 //      CHECK:   %[[T1:.+]] = linalg.init_tensor [8, 33, 4]
 //      CHECK:   %[[T2:.+]] = linalg.generic
 // CHECK-SAME:     indexing_maps = [#[[MAP2]], #[[MAP2]]]
 // CHECK-SAME:     ["parallel", "parallel", "parallel"]
 // CHECK-SAME:     ins(%[[T0]] : tensor<8x33x4xf32>)
 // CHECK-SAME:     outs(%[[T1]] : tensor<8x33x4xf32>)
 //      CHECK:   return %[[T2]] : tensor<8x33x4xf32>
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
 #map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
 func @indexed_consumer_reshape_producer_fusion(%arg0 : tensor<?x?x4x?xi32>,
                                          %arg1 : tensor<?x?x?xi32>) ->
                                          tensor<?x?x?xi32>
 {
-  %0 = linalg.tensor_collapse_shape %arg0 [[0], [1, 2], [3]]:
+  %0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3]]:
     tensor<?x?x4x?xi32> into tensor<?x?x?xi32>
   %1 = linalg.generic {
      indexing_maps = [#map0, #map1, #map1],
      iterator_types = ["parallel", "parallel", "parallel"]}
        ins(%0, %arg1 : tensor<?x?x?xi32>, tensor<?x?x?xi32>)
       outs(%0 : tensor<?x?x?xi32>) {
     ^bb0(%arg3: i32, %arg4: i32, %s: i32):
       %idx0 = linalg.index 0 : index
^bb0(%arg3: i32, %arg4: i32, %s: i32):       // no predecessors
       %idx1 = linalg.index 1 : index
       %1 = arith.muli %arg3, %arg4 : i32
       %2 = arith.index_cast %idx0 : index to i32
       %3 = arith.addi %1, %2 : i32
       %4 = arith.index_cast %idx1 : index to i32
       %5 = arith.addi %3, %4 : i32
       linalg.yield %5 : i32
   } -> tensor<?x?xi32>
-  %1 = linalg.tensor_expand_shape %0 [[0], [1, 2, 3]] :
+  %1 = tensor.expand_shape %0 [[0], [1, 2, 3]] :
     tensor<?x?xi32> into tensor<?x?x4x5xi32>
   return %1 : tensor<?x?x4x5xi32>
 }
 
 // Only check the body in the indexed version of the test.
 //       CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 + d1 * 5 + d2 * 20)>
 //       CHECK: func @indexed_producer_reshape_consumer_fusion
 //       CHECK:   linalg.generic
%c = linalg.generic {
          %2 = arith.index_cast %idx0 : index to i32
          %3 = arith.addi %1, %2 : i32
          %4 = arith.index_cast %idx1 : index to i32
          %5 = arith.addi %3, %4 : i32
          %6 = arith.index_cast %idx2 : index to i32
          %7 = arith.addi %5, %6 : i32
          linalg.yield %7 : i32
        } -> tensor<6x4x210xi32>
-  %d = linalg.tensor_expand_shape %c [[0, 1], [2], [3, 4, 5]]
+  %d = tensor.expand_shape %c [[0, 1], [2], [3, 4, 5]]
        : tensor<6x4x210xi32> into tensor<2x3x4x5x6x7xi32>
   return %d : tensor<2x3x4x5x6x7xi32>
 }
 
 
 //   CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)>
 //   CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)>
 //   CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)>
 //   CHECK-DAG: #[[MAP8:.+]] = affine_map<(d0, d1) -> (d0 + d1 * 3)>
 //   CHECK-DAG: #[[MAP9:.+]] = affine_map<(d0, d1, d2) -> (d0 + d1 * 7 + d2 * 42)>
 //       CHECK: func @reshape_as_consumer_permutation
 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<210x6x4xi32>
 //  CHECK-SAME:   %[[ARG1:.+]]: tensor<210x4xi32>
-//   CHECK-DAG:   %[[T1:.+]] = linalg.tensor_expand_shape %[[ARG0]]
+//   CHECK-DAG:   %[[T1:.+]] = tensor.expand_shape %[[ARG0]]
 //  CHECK-SAME:     [0, 1, 2], [3, 4], [5]
-//   CHECK-DAG:   %[[T2:.+]] = linalg.tensor_expand_shape %[[ARG1]]
+//   CHECK-DAG:   %[[T2:.+]] = tensor.expand_shape %[[ARG1]]
 //  CHECK-SAME:     [0, 1, 2], [3]
 //   CHECK-DAG:   %[[T0:.+]] = linalg.init_tensor [2, 3, 4, 5, 6, 7]
 //       CHECK:   %[[T4:.+]] = linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP5]], #[[MAP6]], #[[MAP7]]]
 //  CHECK-SAME:     ins(%[[T1]], %[[T2]] : tensor<5x6x7x2x3x4xi32>, tensor<5x6x7x4xi32>)
 //  CHECK-SAME:     outs(%[[T0]] : tensor<2x3x4x5x6x7xi32>)
 //       CHECK:   ^{{.+}}(
 //  CHECK-SAME:     %[[ARG8:[a-zA-Z0-9]+]]: i32, %[[ARG9:[a-zA-Z0-9]+]]: i32,

 //       CHECK:       %[[T12:.+]] = arith.index_cast %[[IDX5]]
 //       CHECK:       %[[T13:.+]] = arith.addi %[[T11]], %[[T12]]
 
 // -----
 
 func @reshape_as_producer_projected_permutation(
     %arg0 : tensor<33x8x?xi32>, %shape : tensor<264x?x4xi32>) -> tensor<264x?x4xi32>
 {
-  %0 = linalg.tensor_collapse_shape %arg0 [[0, 1], [2]]
+  %0 = tensor.collapse_shape %arg0 [[0, 1], [2]]
     : tensor<33x8x?xi32> into tensor<264x?xi32>
   %1 = linalg.generic
     {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>,
                       affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
      iterator_types = ["parallel", "parallel", "parallel"]}
      ins(%0 : tensor<264x?xi32>)
     outs(%shape : tensor<264x?x4xi32>) {
   ^bb0(%arg1: i32, %s: i32):  // no predecessors

 //   CHECK-DAG:       %[[T0:.+]] = affine.apply #[[MAP2]](%[[IDX1]], %[[IDX0]])
 //       CHECK:       %[[T1:.+]] = arith.index_cast %[[T0]] : index to i32
 //       CHECK:       %[[T2:.+]] = arith.addi %[[ARG1]], %[[T1]] : i32
 //       CHECK:       %[[T3:.+]] = arith.index_cast %[[IDX2]] : index to i32
 //       CHECK:       %[[T4:.+]] = arith.addi %[[T2]], %[[T3]] : i32
 //       CHECK:       %[[T5:.+]] = arith.index_cast %[[IDX3]] : index to i32
 //       CHECK:       %[[T6:.+]] = arith.addi %[[T4]], %[[T5]] : i32
 //       CHECK:       linalg.yield %[[T6]] : i32
-//       CHECK:    %[[RES2:.+]] = linalg.tensor_collapse_shape %[[RES]]
+//       CHECK:    %[[RES2:.+]] = tensor.collapse_shape %[[RES]]
 //  CHECK-SAME:      [0, 1], [2], [3]
 //  CHECK-SAME:    : tensor<33x8x?x4xi32> into tensor<264x?x4xi32>
 //       CHECK:  return %[[RES2]] : tensor<264x?x4xi32>
 
 // -----
 
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
 #map1 = affine_map<(d0, d1) -> (d1, d0)>
 func @generic_op_reshape_consumer_fusion_projected(%arg0 : tensor<?x?xf32>,
                                                    %arg1 : tensor<?x?xf32>) ->
                                                    tensor<?x?x4x5xf32>
 {
   %0 = linalg.generic {
      indexing_maps = [#map0, #map0, #map1],
      iterator_types = ["parallel", "parallel"]}
        ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
        outs(%arg0 : tensor<?x?xf32>) {
     ^bb0(%arg3: f32, %arg4: f32, %s: f32):       // no predecessors
       %1 = arith.mulf %arg3, %arg4 : f32
       linalg.yield %1 : f32
   } -> tensor<?x?xf32>
-  %1 = linalg.tensor_expand_shape %0 [[0], [1, 2, 3]] :
+  %1 = tensor.expand_shape %0 [[0], [1, 2, 3]] :
     tensor<?x?xf32> into tensor<?x?x4x5xf32>
   return %1 : tensor<?x?x4x5xf32>
 }
 
 //  CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 //  CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)>
 //      CHECK: func @generic_op_reshape_consumer_fusion_projected
 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
 // CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
-//      CHECK:   %[[T0:.+]] = linalg.tensor_expand_shape %[[ARG0]]
+//      CHECK:   %[[T0:.+]] = tensor.expand_shape %[[ARG0]]
 // CHECK-SAME:     [0, 1, 2], [3]
 // CHECK-SAME:     tensor<?x?xf32> into tensor<?x4x5x?xf32>
-//      CHECK:   %[[T1:.+]] = linalg.tensor_expand_shape %[[ARG1]]
+//      CHECK:   %[[T1:.+]] = tensor.expand_shape %[[ARG1]]
 // CHECK-SAME:     [0, 1, 2], [3]
 // CHECK-SAME:     tensor<?x?xf32> into tensor<?x4x5x?xf32>
 //      CHECK:   %[[T3:.+]] = linalg.generic
 // CHECK-SAME:     indexing_maps = [#[[MAP4]], #[[MAP4]], #[[MAP5]]]
 // CHECK-SAME:     ["parallel", "parallel", "parallel", "parallel"]
 // CHECK-SAME:     ins(%[[T0]], %[[T1]] : tensor<?x4x5x?xf32>, tensor<?x4x5x?xf32>)
 // CHECK-SAME:     outs(%{{.+}} : tensor<?x?x4x5xf32>)
 //      CHECK:   return %[[T3]] : tensor<?x?x4x5xf32>
 
 // -----
 
 func @unit_dim_reshape_expansion(%arg0 : tensor<1x5xf32>) -> tensor<5x5xf32> {
-  %0 = linalg.tensor_collapse_shape %arg0 [[0, 1]]
+  %0 = tensor.collapse_shape %arg0 [[0, 1]]
       : tensor<1x5xf32> into tensor<5xf32>
   %1 = linalg.init_tensor [5, 5] : tensor<5x5xf32>
   %2 = linalg.generic
     {indexing_maps = [affine_map<(d0, d1) -> (d0)>,
                       affine_map<(d0, d1) -> (d0, d1)>],
      iterator_types = ["parallel", "parallel"]}
     ins(%0 : tensor<5xf32>) outs(%1 : tensor<5x5xf32>) {
   ^bb0(%arg2: f32, %arg3: f32):  // no predecessors
     linalg.yield %arg2 : f32
   } -> tensor<5x5xf32>
   return %2 : tensor<5x5xf32>
 }
 //      CHECK: func @unit_dim_reshape_expansion
-//  CHECK-DAG:   linalg.tensor_collapse_shape
+//  CHECK-DAG:   tensor.collapse_shape
 //  CHECK-DAG:   linalg.init_tensor
 //      CHECK:   linalg.generic
 
 // -----
 
 func @unit_dim_reshape_collapse(%arg0 : tensor<5xf32>) -> tensor<5x1x5xf32> {
   %0 = linalg.init_tensor [5, 5] : tensor<5x5xf32>
   %1 = linalg.generic
     {indexing_maps = [affine_map<(d0, d1) -> (d0)>,
                       affine_map<(d0, d1) -> (d0, d1)>],
      iterator_types = ["parallel", "parallel"]}
     ins(%arg0 : tensor<5xf32>) outs(%0 : tensor<5x5xf32>) {
   ^bb0(%arg2: f32, %arg3: f32):  // no predecessors
     linalg.yield %arg2 : f32
   } -> tensor<5x5xf32>
-  %2 = linalg.tensor_expand_shape %1 [[0, 1], [2]]
+  %2 = tensor.expand_shape %1 [[0, 1], [2]]
     : tensor<5x5xf32> into tensor<5x1x5xf32>
   return %2 : tensor<5x1x5xf32>
 }
 // CHECK: func @unit_dim_reshape_collapse
 // CHECK:   linalg.init_tensor
 // CHECK:   linalg.generic
-// CHECK:   linalg.tensor_expand_shape
+// CHECK:   tensor.expand_shape
 
 // -----
 
 func @unit_dim_reshape_expansion_full
   (%arg0 : tensor<1x?x1x2x1x4xf32>, %arg1 : tensor<?x2x4xf32>)
   -> tensor<?x2x4xf32> {
   %c1 = arith.constant 1 : index
-  %0 = linalg.tensor_collapse_shape %arg0 [[0, 1, 2], [3, 4], [5]]
+  %0 = tensor.collapse_shape %arg0 [[0, 1, 2], [3, 4], [5]]
     : tensor<1x?x1x2x1x4xf32> into tensor<?x2x4xf32>
   %1 = tensor.dim %arg0, %c1 : tensor<1x?x1x2x1x4xf32>
   %2 = linalg.init_tensor [%1, 2, 4] : tensor<?x2x4xf32>
   %3 = linalg.generic
     {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                       affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                       affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
      iterator_types = ["parallel", "parallel", "parallel"]}
     ins(%0, %arg1 : tensor<?x2x4xf32>, tensor<?x2x4xf32>)
     outs(%2 : tensor<?x2x4xf32>) {
   ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):  // no predecessors
     %4 = arith.mulf %arg2, %arg3 : f32
     linalg.yield %4 : f32
   } -> tensor<?x2x4xf32>
   return %3 : tensor<?x2x4xf32>
 }
 //      CHECK: func @unit_dim_reshape_expansion_full
-//  CHECK-DAG:   linalg.tensor_collapse_shape
+//  CHECK-DAG:   tensor.collapse_shape
 //  CHECK-DAG:   linalg.init_tensor
 //      CHECK:   linalg.generic
 // CHECK-SAME:     ins(%{{.+}}, %{{.+}} : tensor<?x2x4xf32>, tensor<?x2x4xf32>)
 
 //         FOLDUNITDIM: func @unit_dim_reshape_expansion_full
 //    FOLDUNITDIM-SAME:   %[[ARG0:.+]]: tensor<1x?x1x2x1x4xf32>
 //    FOLDUNITDIM-SAME:   %[[ARG1:.+]]: tensor<?x2x4xf32>
-//     FOLDUNITDIM-DAG:   %[[RESHAPE:.+]] = linalg.tensor_expand_shape %[[ARG1]]
+//     FOLDUNITDIM-DAG:   %[[RESHAPE:.+]] = tensor.expand_shape %[[ARG1]]
 //         FOLDUNITDIM:   linalg.generic
 //    FOLDUNITDIM-SAME:     ins(%[[ARG0]], %[[RESHAPE]] : tensor<1x?x1x2x1x4xf32>, tensor<1x?x1x2x1x4xf32>)
 //    FOLDUNITDIM-SAME:     outs(%{{.+}} : tensor<1x?x1x2x1x4xf32>)
 

mlir/test/Dialect/Linalg/reshape_linearization_fusion.mlir

 // RUN: mlir-opt -split-input-file -linalg-fold-reshape-ops-by-linearization %s | FileCheck %s
 
 #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 func @generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x?xi32>)
   -> tensor<?x?x4x?xi32> {
-  %0 = linalg.tensor_expand_shape %arg0 [[0], [1, 2], [3]] :
+  %0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] :
     tensor<?x?x?xi32> into tensor<?x?x4x?xi32>
   %1 = linalg.generic {
     indexing_maps = [#map0, #map0],
     iterator_types = ["parallel", "parallel", "parallel", "parallel"] }
     ins(%0 : tensor<?x?x4x?xi32>)
     outs(%0 : tensor<?x?x4x?xi32>) {
   ^bb0(%arg6: i32, %arg7 : i32):       // no predecessors
     %idx = linalg.index 0 : index
     %2 = arith.index_cast %idx : index to i32
     %3 = arith.addi %arg6, %2 : i32
     linalg.yield %3 : i32
   } -> tensor<?x?x4x?xi32>
   return %1 : tensor<?x?x4x?xi32>
 }
 //   CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 4 + d2, d3)>
 //   CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 //       CHECK: func @generic_op_reshape_producer_fusion
 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<?x?x?xi32>
-//       CHECK:   %[[T0:.+]] = linalg.tensor_expand_shape %[[ARG0]]
+//       CHECK:   %[[T0:.+]] = tensor.expand_shape %[[ARG0]]
 //  CHECK-SAME:     [0], [1, 2], [3]
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP3]], #[[MAP4]]]
 //  CHECK-SAME:     ins(%[[ARG0]] : tensor<?x?x?xi32>)
 //  CHECK-SAME:     outs(%[[T0]] : tensor<?x?x4x?xi32>)
 //       CHECK:   %[[IDX:.+]] = linalg.index 0 : index
 //  CHECK-NEXT:   %[[IDX_CASTED:.+]] = arith.index_cast %[[IDX]] : index to i32
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 func @generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?x4x5xi32>)
   -> tensor<?x?xi32> {
   %0 = linalg.generic {
     indexing_maps = [#map0, #map0],
     iterator_types = ["parallel", "parallel", "parallel", "parallel"] }
     ins(%arg0 : tensor<?x?x4x5xi32>) outs(%arg0 : tensor<?x?x4x5xi32>) {
   ^bb0(%arg6: i32, %arg7: i32):       // no predecessors
     %idx = linalg.index 0 : index
     %2 = arith.index_cast %idx : index to i32
     %3 = arith.addi %arg6, %2 : i32
     linalg.yield %3 : i32
   } -> tensor<?x?x4x5xi32>
-  %1 = linalg.tensor_collapse_shape %0 [[0], [1, 2, 3]] :
+  %1 = tensor.collapse_shape %0 [[0], [1, 2, 3]] :
     tensor<?x?x4x5xi32> into tensor<?x?xi32>
   return %1 : tensor<?x?xi32>
 }
 //   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 //   CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>
 //       CHECK: func @generic_op_reshape_consumer_fusion
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x4x5xi32>
-//       CHECK:   %[[T0:.+]] = linalg.tensor_collapse_shape %[[ARG0]]
+//       CHECK:   %[[T0:.+]] = tensor.collapse_shape %[[ARG0]]
 //  CHECK-SAME:     [0], [1, 2, 3]
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP2]], #[[MAP3]]]
 //  CHECK-SAME:     outs(%[[T0]] : tensor<?x?xi32>)
 //       CHECK:   %[[IDX:.+]] = linalg.index 0 : index
 //  CHECK-NEXT:   %[[IDX_CASTED:.+]] = arith.index_cast %[[IDX]] : index to i32
-//   CHECK-NOT:   linalg.tensor_collapse_shape
+//   CHECK-NOT:   tensor.collapse_shape
 
 // -----
 
 #map2 = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
 #map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 func @generic_op_021_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<3x7x5xf32> {
-  %0 = linalg.tensor_expand_shape %arg0 [[0], [1, 2]]
+  %0 = tensor.expand_shape %arg0 [[0], [1, 2]]
       : tensor<3x35xf32> into tensor<3x5x7xf32>
   %1 = linalg.init_tensor [3, 7, 5] : tensor<3x7x5xf32>
   %2 = linalg.generic
     {indexing_maps = [#map2, #map3],
      iterator_types = ["parallel", "parallel", "parallel"]}
     ins(%0 : tensor<3x5x7xf32>) outs(%1 : tensor<3x7x5xf32>) {
     ^bb0(%arg2: f32, %arg3 : f32):  // no predecessors
       linalg.yield %arg2 : f32
     } -> tensor<3x7x5xf32>
     return %2 : tensor<3x7x5xf32>
 }
 
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 + d2 * 7)>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 //       CHECK: func @generic_op_021_permultation_reshape_producer_fusion
-//   CHECK-NOT:   linalg.tensor_expand_shape
+//   CHECK-NOT:   tensor.expand_shape
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
 
 // -----
 
 #map2 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
 #map3 = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
 func @generic_op_120_permutation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<5x7x3xf32> {
-  %0 = linalg.tensor_expand_shape %arg0 [[0], [1, 2]]
+  %0 = tensor.expand_shape %arg0 [[0], [1, 2]]
       : tensor<3x35xf32> into tensor<3x5x7xf32>
   %1 = linalg.init_tensor [5, 7, 3] : tensor<5x7x3xf32>
   %2 = linalg.generic
     {indexing_maps = [#map2, #map3],
      iterator_types = ["parallel", "parallel", "parallel"]}
     ins(%0 : tensor<3x5x7xf32>) outs(%1 : tensor<5x7x3xf32>) {
     ^bb0(%arg2: f32, %arg3: f32):  // no predecessors
       linalg.yield %arg2 : f32
     } -> tensor<5x7x3xf32>
     return %2 : tensor<5x7x3xf32>
 }
 
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d1, d0 * 7 + d2)>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
 //       CHECK: func @generic_op_120_permutation_reshape_producer_fusion
-//   CHECK-NOT:   linalg.tensor_expand_shape
+//   CHECK-NOT:   tensor.expand_shape
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2) -> (d0)>
 #map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
 #map2 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
 #map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 func @generic_op_102_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<5x3x7xf32> {
-  %0 = linalg.tensor_expand_shape %arg0 [[0], [1, 2]]
+  %0 = tensor.expand_shape %arg0 [[0], [1, 2]]
       : tensor<3x35xf32> into tensor<3x5x7xf32>
   %1 = linalg.init_tensor [5, 3, 7] : tensor<5x3x7xf32>
   %2 = linalg.generic
     {indexing_maps = [#map2, #map3],
      iterator_types = ["parallel", "parallel", "parallel"]}
     ins(%0 : tensor<3x5x7xf32>) outs(%1 : tensor<5x3x7xf32>) {
     ^bb0(%arg2: f32, %arg3: f32):  // no predecessors
       linalg.yield %arg2 : f32
     } -> tensor<5x3x7xf32>
     return %2 : tensor<5x3x7xf32>
 }
 
 
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d1, d0 * 7 + d2)>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 //       CHECK: func @generic_op_102_permultation_reshape_producer_fusion
-//   CHECK-NOT:   linalg.tensor_expand_shape
+//   CHECK-NOT:   tensor.expand_shape
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 #map1 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
 #map2 = affine_map<(d0, d1, d2) -> (d0)>
 #map3 = affine_map<(d0, d1, d2) -> (d1, d2)>
 func @generic_op_102_permultation_reshape_consumer_fusion(%arg0 : tensor<3x5x7xf32>) -> tensor<5x21xf32> {
   %0 = linalg.init_tensor [5, 3, 7] : tensor<5x3x7xf32>
   %1 = linalg.generic
     {indexing_maps = [#map0, #map1],
      iterator_types = ["parallel", "parallel", "parallel"]}
     ins(%arg0 : tensor<3x5x7xf32>) outs(%0 : tensor<5x3x7xf32>) {
     ^bb0(%arg2: f32, %arg3 : f32):  // no predecessors
       linalg.yield %arg2 : f32
   } -> tensor<5x3x7xf32>
-  %2 = linalg.tensor_collapse_shape %1 [[0], [1, 2]]
+  %2 = tensor.collapse_shape %1 [[0], [1, 2]]
       : tensor<5x3x7xf32> into tensor<5x21xf32>
   return %2 : tensor<5x21xf32>
 }
 //   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 //   CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d1, d0 * 7 + d2)>
 //       CHECK: func @generic_op_102_permultation_reshape_consumer_fusion
 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<3x5x7xf32>
 //       CHECK:   %[[T0:.+]] = linalg.init_tensor [5, 3, 7]
-//       CHECK:   %[[T1:.+]] = linalg.tensor_collapse_shape %[[T0]]
+//       CHECK:   %[[T1:.+]] = tensor.collapse_shape %[[T0]]
 //  CHECK-SAME:     [0], [1, 2]
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP2]], #[[MAP3]]]
 //  CHECK-SAME:     ins(%[[ARG0]] : tensor<3x5x7xf32>)
 //  CHECK-SAME:     outs(%[[T1]] : tensor<5x21xf32>)
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 func @generic_op_reshape_consumer_nofusion(%arg0 : tensor<?x?x?x5xf32>,
                                            %arg1 : tensor<?x?x?x5xf32>) ->
                                            tensor<?x?xf32>
 {
   %0 = linalg.generic {
      indexing_maps = [#map0, #map0, #map0],
      iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
       ins(%arg0, %arg1 : tensor<?x?x?x5xf32>, tensor<?x?x?x5xf32>)
       outs(%arg0 : tensor<?x?x?x5xf32>) {
     ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):       // no predecessors
       %1 = arith.mulf %arg3, %arg4 : f32
       linalg.yield %1 : f32
   } -> tensor<?x?x?x5xf32>
-  %1 = linalg.tensor_collapse_shape %0 [[0], [1, 2, 3]] :
+  %1 = tensor.collapse_shape %0 [[0], [1, 2, 3]] :
     tensor<?x?x?x5xf32> into tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
 // CHECK-LABEL: func @generic_op_reshape_consumer_nofusion
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?x5xf32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?x5xf32>
 //       CHECK:   %[[NOFUSE:.+]] = linalg.generic
 //  CHECK-SAME:     ins(%[[ARG0]], %[[ARG1]]
-//       CHECK:   %[[RESULT:.+]] = linalg.tensor_collapse_shape %[[NOFUSE]]
+//       CHECK:   %[[RESULT:.+]] = tensor.collapse_shape %[[NOFUSE]]
 //       CHECK:   return %[[RESULT]]
 
 
 // -----
 
 func @generic_op_permultation_reshape_consumer_fusion_unused_dim(%arg0 : tensor<6x1xf32>) -> tensor<6xi32> {
   %0 = linalg.init_tensor [6, 1] : tensor<6x1xi32>
   %1 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
                                         affine_map<(d0, d1) -> (d0, d1)>],
    iterator_types = ["parallel", "parallel"]}
    ins(%arg0 : tensor<6x1xf32>) outs(%0 : tensor<6x1xi32>) {
     ^bb0(%arg3: f32, %arg4: i32):  // no predecessors
       %5 = arith.fptosi %arg3 : f32 to i32
       linalg.yield %5 : i32
     } -> tensor<6x1xi32>
-    %6 = linalg.tensor_collapse_shape %1 [[0, 1]] : tensor<6x1xi32> into tensor<6xi32>
+    %6 = tensor.collapse_shape %1 [[0, 1]] : tensor<6x1xi32> into tensor<6xi32>
   return %6 : tensor<6xi32>
 }
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0)>
 //       CHECK: func @generic_op_permultation_reshape_consumer_fusion_unused_dim
 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<6x1xf32>
 //       CHECK:   %[[T0:.+]] = linalg.init_tensor [6, 1]
-//       CHECK:   %[[T1:.+]] = linalg.tensor_collapse_shape %[[T0]]
+//       CHECK:   %[[T1:.+]] = tensor.collapse_shape %[[T0]]
 //  CHECK-SAME:   [0, 1]
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
 //  CHECK-SAME:     ins(%[[ARG0]] : tensor<6x1xf32>)
 //  CHECK-SAME:     outs(%[[T1]] : tensor<6xi32>)

mlir/test/Dialect/Linalg/reshape_linearization_fusion_with_unit_dims.mlir

%3 = linalg.generic {
       indexing_maps = [#map, #map, #map],
       iterator_types = ["parallel", "parallel"]}
       ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
       outs(%2 : tensor<?x?xf32>) {
       ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
         %4 = arith.addf %arg2, %arg3 : f32
         linalg.yield %4 : f32
       } -> tensor<?x?xf32>
-  %4 = linalg.tensor_expand_shape %3 [[0], [1, 2]] : tensor<?x?xf32> into tensor<?x?x1xf32>
+  %4 = tensor.expand_shape %3 [[0], [1, 2]] : tensor<?x?xf32> into tensor<?x?x1xf32>
   return %4 : tensor<?x?x1xf32>
 }
 // CHECK-LABEL: func @do_not_fold1
 //       CHECK: %[[VAL:.+]] = linalg.generic
-//       CHECK: linalg.tensor_expand_shape %[[VAL]]
+//       CHECK: tensor.expand_shape %[[VAL]]
 
 // -----
 
 #map = affine_map<(d0, d1) -> (d0, d1)>
 func @do_not_fold2(%arg0 : tensor<?x?x1xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32>
 {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
-  %0 = linalg.tensor_collapse_shape %arg0 [[0], [1, 2]] : tensor<?x?x1xf32> into tensor<?x?xf32>
+  %0 = tensor.collapse_shape %arg0 [[0], [1, 2]] : tensor<?x?x1xf32> into tensor<?x?xf32>
   %1 = tensor.dim %arg1, %c0 : tensor<?x?xf32>
   %2 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
   %3 = linalg.init_tensor [%1, %2] : tensor<?x?xf32>
   %4 = linalg.generic {
       indexing_maps = [#map, #map, #map],
       iterator_types = ["parallel", "parallel"]}
       ins(%0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
       outs(%3 : tensor<?x?xf32>) {
       ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
         %4 = arith.addf %arg2, %arg3 : f32
         linalg.yield %4 : f32
       } -> tensor<?x?xf32>
   return %4 : tensor<?x?xf32>
 }
 // CHECK-LABEL: func @do_not_fold2
-//       CHECK: %[[VAL:.+]] = linalg.tensor_collapse_shape
+//       CHECK: %[[VAL:.+]] = tensor.collapse_shape
 //       CHECK: linalg.generic
 //  CHECK-SAME:   ins(%[[VAL]], %{{.+}} : tensor<?x?xf32>, tensor<?x?xf32>)

mlir/test/Dialect/Linalg/resolve-shaped-type-result-dims.mlir

 
 // -----
 
 func @dim_reshape_expansion(%arg0 : tensor<6x5x?xf32>) -> (index, index, index)
 {
   %c1 = arith.constant 1 : index
   %c3 = arith.constant 3 : index
   %c4 = arith.constant 4 : index
-  %0 = linalg.tensor_expand_shape %arg0 [[0, 1], [2], [3, 4, 5]]
+  %0 = tensor.expand_shape %arg0 [[0, 1], [2], [3, 4, 5]]
       : tensor<6x5x?xf32> into tensor<2x3x5x4x?x7xf32>
   %1 = tensor.dim %0, %c1 : tensor<2x3x5x4x?x7xf32>
   %2 = tensor.dim %0, %c3 : tensor<2x3x5x4x?x7xf32>
   %3 = tensor.dim %0, %c4 : tensor<2x3x5x4x?x7xf32>
   return %1, %2, %3 : index, index, index
 }
 //      CHECK: #[[MAP:.+]] = affine_map<()[s0] -> (s0 floordiv 28)>
 //      CHECK: func @dim_reshape_expansion
 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<6x5x?xf32>
 //  CHECK-DAG:   %[[C2:.+]] = arith.constant 2 : index
 //  CHECK-DAG:   %[[C3:.+]] = arith.constant 3 : index
 //  CHECK-DAG:   %[[C4:.+]] = arith.constant 4 : index
 //      CHECK:   %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C2]]
 //      CHECK:   %[[D1:.+]] = affine.apply #[[MAP]]()[%[[D0]]]
 //      CHECK:   return %[[C3]], %[[C4]], %[[D1]]
 
 // -----
 
 func @dim_reshape_collapse(%arg0 : tensor<2x3x5x4x?x7xf32>) -> (index, index)
 {
   %c1 = arith.constant 1 : index
   %c2 = arith.constant 2 : index
-  %0 = linalg.tensor_collapse_shape %arg0 [[0, 1], [2], [3, 4, 5]]
+  %0 = tensor.collapse_shape %arg0 [[0, 1], [2], [3, 4, 5]]
       : tensor<2x3x5x4x?x7xf32> into tensor<6x5x?xf32>
   %1 = tensor.dim %0, %c1 : tensor<6x5x?xf32>
   %2 = tensor.dim %0, %c2 : tensor<6x5x?xf32>
   return %1, %2 : index, index
 }
 //      CHECK: #[[MAP:.+]] = affine_map<()[s0] -> (s0 * 28)>
 //      CHECK: func @dim_reshape_collapse
 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<2x3x5x4x?x7xf32>

mlir/test/Dialect/Linalg/roundtrip.mlir

 // CHECK-LABEL: func @named_ops
 //       CHECK:   linalg.batch_matmul
 //       CHECK:   linalg.batch_matmul
 //       CHECK:   linalg.batch_matmul
 //       CHECK:   linalg.batch_matmul
 
 // -----
 
-func @tensor_reshape_zero_dim(%arg0 : tensor<1x1xf32>, %arg1 : tensor<f32>) -> (tensor<f32>, tensor<1x1xf32>)
-{
-  %0 = linalg.tensor_collapse_shape %arg0 [] : tensor<1x1xf32> into tensor<f32>
-  %1 = linalg.tensor_expand_shape %0 [] : tensor<f32> into tensor<1x1xf32>
-  return %0, %1 : tensor<f32>, tensor<1x1xf32>
-}
-// CHECK-LABEL: func @tensor_reshape_zero_dim
-//       CHECK:   linalg.tensor_collapse_shape %{{.*}} [] : tensor<1x1xf32> into tensor<f32>
-//       CHECK:   linalg.tensor_expand_shape %{{.*}} [] : tensor<f32> into tensor<1x1xf32>
-
-// -----
-
-
 func @init_tensor(%arg0 : index, %arg1 : index)
 {
   %0 = linalg.init_tensor [3, 42] : tensor<3x42xf32>
   %1 = linalg.init_tensor [4, %arg0, %arg1, 5] : tensor<4x?x?x5xf32>
   return
 }
 // CHECK-LABEL: func @init_tensor
 //       CHECK:   linalg.init_tensor [3, 42] : tensor<3x42xf32>
 //       CHECK:   linalg.init_tensor [4, %{{.*}}, %{{.*}}, 5] : tensor<4x?x?x5xf32>
 
 // -----
 
-func @legal_collapsing_reshape_dynamic_tensor
-  (%arg0: tensor<?x?x?x4x?xf32>) -> tensor<?x?x?xf32>
-{
-  %0 = linalg.tensor_collapse_shape %arg0 [[0], [1], [2, 3, 4]] :
-    tensor<?x?x?x4x?xf32> into tensor<?x?x?xf32>
-  return %0 : tensor<?x?x?xf32>
-}
-//      CHECK: func @legal_collapsing_reshape_dynamic_tensor
-//      CHECK:   linalg.tensor_collapse_shape
-// CHECK-SAME:    [0], [1], [2, 3, 4]
-
-// -----
-
 func @fill_tensor(%arg0 : index, %arg1 : index, %arg2 : f32) -> tensor<?x?xf32> {
   %0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32>
   %1 = linalg.fill(%arg2, %0) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
 // CHECK: %{{.+}} = linalg.fill(%{{.+}}, %{{.+}}) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
 
 // -----

mlir/test/Dialect/MemRef/ops.mlir

func @expand_collapse_shape_static(%arg0: memref<3x4x5xf32>,
   %r2 = memref.expand_shape %2 [[0, 1, 2]] :
     memref<60xf32> into memref<3x4x5xf32>
   // Reshapes that expand and collapse back a contiguous buffer with some 1's.
   %3 = memref.expand_shape %arg0 [[0, 1], [2], [3, 4]] :
     memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
   %r3 = memref.collapse_shape %3 [[0, 1], [2], [3, 4]] :
     memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
   // Reshapes on tensors.
-  %t0 = linalg.tensor_expand_shape %arg1 [[0, 1], [2], [3, 4]] :
+  %t0 = tensor.expand_shape %arg1 [[0, 1], [2], [3, 4]] :
     tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>
-  %rt0 = linalg.tensor_collapse_shape %t0 [[0, 1], [2], [3, 4]] :
+  %rt0 = tensor.collapse_shape %t0 [[0, 1], [2], [3, 4]] :
     tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>
-  %t1 = linalg.tensor_expand_shape %arg2 [[0, 1], [2], [3, 4]] :
+  %t1 = tensor.expand_shape %arg2 [[0, 1], [2], [3, 4]] :
     tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>
-  %rt1 = linalg.tensor_collapse_shape %t1 [[0], [1, 2], [3, 4]] :
+  %rt1 = tensor.collapse_shape %t1 [[0], [1, 2], [3, 4]] :
     tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>
   return
 }
 // CHECK-LABEL: func @expand_collapse_shape_static
 //       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]
 //  CHECK-SAME:     memref<3x4x5xf32> into memref<12x5xf32>
 //       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
 //  CHECK-SAME:     memref<12x5xf32> into memref<3x4x5xf32>
 //       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0], [1, 2]]
 //  CHECK-SAME:     memref<3x4x5xf32> into memref<3x20xf32>
 //       CHECK:   memref.expand_shape {{.*}} {{\[}}[0], [1, 2]]
 //  CHECK-SAME:     memref<3x20xf32> into memref<3x4x5xf32>
 //       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1, 2]]
 //  CHECK-SAME:     memref<3x4x5xf32> into memref<60xf32>
 //       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1, 2]]
 //  CHECK-SAME:     memref<60xf32> into memref<3x4x5xf32>
 //       CHECK:   memref.expand_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]
 //  CHECK-SAME:     memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
 //       CHECK:   memref.collapse_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]
 //  CHECK-SAME:     memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
 //
-//       CHECK:   linalg.tensor_expand_shape {{.*}}: tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>
-//       CHECK:   linalg.tensor_collapse_shape {{.*}}: tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>
-//       CHECK:   linalg.tensor_expand_shape {{.*}}: tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>
-//       CHECK:   linalg.tensor_collapse_shape {{.*}}: tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>
+//       CHECK:   tensor.expand_shape {{.*}}: tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>
+//       CHECK:   tensor.collapse_shape {{.*}}: tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>
+//       CHECK:   tensor.expand_shape {{.*}}: tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>
+//       CHECK:   tensor.collapse_shape {{.*}}: tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>
 
 
 func @expand_collapse_shape_dynamic(%arg0: memref<?x?x?xf32>,
          %arg1: memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]>,
          %arg2: memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]>,
          %arg3: memref<?x42xf32, offset : 0, strides : [42, 1]>) {
   %0 = memref.collapse_shape %arg0 [[0, 1], [2]] :
     memref<?x?x?xf32> into memref<?x?xf32>

mlir/test/Dialect/Tensor/canonicalize.mlir

func @fold_overlapping_insert(%input : tensor<?x?x?xf32>, %slice1: tensor<4x?x8xf32>, %slice2: tensor<4x?x8xf32>, %i: index, %size: index) -> (tensor<?x?x?xf32>) {
   %c0 = arith.constant 0: index
   %c1 = arith.constant 1: index
   %0 = tensor.insert_slice %slice1 into %input[%c0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<4x?x8xf32> into tensor<?x?x?xf32>
   // CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[SLICE2]] into %[[INPUT]]
   %1 = tensor.insert_slice %slice2 into %0[%c0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<4x?x8xf32> into tensor<?x?x?xf32>
   // CHECK: return %[[INSERT]]
   return %1 : tensor<?x?x?xf32>
 }
+
+// -----
+
+func @expanding_tensor_reshapes(%arg0 : tensor<?x?xf32>)
+    -> tensor<?x6x4x?x5xf32> {
+  %0 = tensor.expand_shape %arg0 [[0, 1], [2]]
+      : tensor<?x?xf32> into tensor<?x4x?xf32>
+  %1 = tensor.expand_shape %0 [[0, 1], [2], [3, 4]]
+      : tensor<?x4x?xf32> into tensor<?x6x4x?x5xf32>
+  return %1 : tensor<?x6x4x?x5xf32>
+}
+// CHECK-LABEL: expanding_tensor_reshapes
+//       CHECK:   tensor.expand_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
+//   CHECK-NOT:   tensor.expand_shape
+
+// -----
+
+func @expanding_tensor_reshapes_to_zero_dim(%arg0 : tensor<f32>)
+    -> tensor<1x1x1xf32> {
+  %0 = tensor.expand_shape %arg0 [] : tensor<f32> into tensor<1xf32>
+  %1 = tensor.expand_shape %0 [[0, 1, 2]]
+      : tensor<1xf32> into tensor<1x1x1xf32>
+  return %1 : tensor<1x1x1xf32>
+}
+// CHECK-LABEL: expanding_tensor_reshapes_to_zero
+//       CHECK:   tensor.expand_shape %{{.*}} []
+//  CHECK-SAME:     tensor<f32> into tensor<1x1x1xf32>
+
+// -----
+
+func @fold_tensor_reshape(%arg0 : tensor<12x4xf32>) -> tensor<12x4xf32> {
+  %0 = tensor.expand_shape %arg0 [[0, 1], [2]]
+      : tensor<12x4xf32> into tensor<3x4x4xf32>
+  %1 = tensor.collapse_shape %0 [[0, 1], [2]]
+      : tensor<3x4x4xf32> into tensor<12x4xf32>
+  return %1 : tensor<12x4xf32>
+}
+// CHECK-LABEL: @fold_tensor_reshape
+//   CHECK-NOT:   linalg.{{.*}}shape
+
+// -----
+
+func @fold_tensor_reshape_dynamic(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %0 = tensor.expand_shape %arg0 [[0, 1], [2]]
+      : tensor<?x?xf32> into tensor<?x4x?xf32>
+  %1 = tensor.collapse_shape %0 [[0, 1], [2]]
+      : tensor<?x4x?xf32> into tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+// CHECK-LABEL: @fold_tensor_reshape_dynamic
+//   CHECK-NOT:   linalg.{{.*}}_shape
+
+// -----
+func @reshape_collapse(%arg0 : tensor<2x3x4x5x6x7x8xf32>)
+    -> tensor<24x5x42x8xf32> {
+  %0 = tensor.collapse_shape %arg0 [[0, 1, 2, 3, 4, 5, 6]]
+      : tensor<2x3x4x5x6x7x8xf32> into tensor<40320xf32>
+  %1 = tensor.expand_shape %0 [[0, 1, 2, 3]]
+      : tensor<40320xf32> into tensor<24x5x42x8xf32>
+  return %1 : tensor<24x5x42x8xf32>
+}
+//      CHECK: func @reshape_collapse
+// CHECK-SAME:   %[[ARG0:.+]]: tensor<2x3x4x5x6x7x8xf32>
+//      CHECK:   %[[RESULT:.+]] = tensor.collapse_shape %[[ARG0]]
+// CHECK-SAME:     [0, 1, 2], [3], [4, 5], [6]
+//      CHECK:   return %[[RESULT]]
+
+// -----
+
+func @reshape_expand(%arg0 : tensor<24x5x42x8xf32>)
+    -> tensor<2x3x4x5x6x7x8xf32> {
+  %0 = tensor.collapse_shape %arg0 [[0, 1, 2, 3]]
+      : tensor<24x5x42x8xf32> into tensor<40320xf32>
+  %1 = tensor.expand_shape %0 [[0, 1, 2, 3, 4, 5, 6]]
+      : tensor<40320xf32> into tensor<2x3x4x5x6x7x8xf32>
+  return %1 : tensor<2x3x4x5x6x7x8xf32>
+}
+//      CHECK: func @reshape_expand
+// CHECK-SAME:   %[[ARG0:.+]]: tensor<24x5x42x8xf32>
+//      CHECK:   %[[RESULT:.+]] = tensor.expand_shape %[[ARG0]]
+// CHECK-SAME:     [0, 1, 2], [3], [4, 5], [6]
+//      CHECK:   return %[[RESULT]]
+
+// -----
+
+func @expand_reshape_1D(%arg0 : tensor<2048xf32>) -> tensor<4x512xf32> {
+  %0 = tensor.expand_shape %arg0 [[0, 1, 2, 3]]
+    : tensor<2048xf32> into tensor<1x4x1x512xf32>
+  %1 = tensor.collapse_shape %0 [[0, 1, 2], [3]]
+    : tensor<1x4x1x512xf32> into tensor<4x512xf32>
+  return %1 : tensor<4x512xf32>
+}
+//       CHECK: func @expand_reshape_1D
+//       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1]]
+//  CHECK-SAME:   tensor<2048xf32> into tensor<4x512xf32>
+
+// -----
+
+// CHECK-LABEL: zero_rank_reshape_multi
+func @zero_rank_reshape_multi(%arg0: tensor<f32>) -> tensor<f32> {
+  // CHECK: return %arg0
+  %0 = tensor.expand_shape %arg0 [] : tensor<f32> into tensor<1xf32>
+  %1 = tensor.expand_shape %0 [[0, 1]] : tensor<1xf32> into tensor<1x1xf32>
+  %2 = tensor.collapse_shape %1 [] : tensor<1x1xf32> into tensor<f32>
+  return %2 : tensor<f32>
+}
+
+// -----
+
+func @collapsing_tensor_reshapes(%arg0 : tensor<?x?x?x?x?xf32>)
+    -> tensor<?x?xf32> {
+  %0 = tensor.collapse_shape %arg0 [[0, 1], [2], [3, 4]]
+      : tensor<?x?x?x?x?xf32> into tensor<?x?x?xf32>
+  %1 = tensor.collapse_shape %0 [[0, 1], [2]]
+      : tensor<?x?x?xf32> into tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+// CHECK-LABEL: collapsing_tensor_reshapes
+//       CHECK:   tensor.collapse_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
+//   CHECK-NOT:   tensor.collapse_shape
+
+// -----
+
+func @collapsing_tensor_reshapes_to_zero_dim(%arg0 : tensor<1x1x1xf32>)
+    -> tensor<f32> {
+  %0 = tensor.collapse_shape %arg0 [[0, 1, 2]]
+      : tensor<1x1x1xf32> into tensor<1xf32>
+  %1 = tensor.collapse_shape %0 [] : tensor<1xf32> into tensor<f32>
+  return %1 : tensor<f32>
+}
+// CHECK-LABEL: collapsing_tensor_reshapes_to_zero
+//       CHECK:   tensor.collapse_shape %{{.*}} []
+//  CHECK-SAME:     tensor<1x1x1xf32> into tensor<f32>
+
+// -----
+
+func @fold_reshape_1D(%arg0 : tensor<4x512xf32>) -> tensor<2048xf32> {
+  %0 = tensor.expand_shape %arg0 [[0, 1, 2], [3]]
+    : tensor<4x512xf32> into tensor<1x4x1x512xf32>
+  %1 = tensor.collapse_shape %0 [[0, 1, 2, 3]]
+    : tensor<1x4x1x512xf32> into tensor<2048xf32>
+  return %1 : tensor<2048xf32>
+}
+//       CHECK: func @fold_reshape_1D
+//       CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1]]
+//  CHECK-SAME:   tensor<4x512xf32> into tensor<2048xf32>
+
+// -----
+
+func @fold_reshape_unit_dims(%arg0 : tensor<2048x1x1xf32>)
+    -> tensor<4x512x1x1xf32> {
+  %0 = tensor.expand_shape %arg0 [[0, 1, 2, 3], [4], [5]]
+    : tensor<2048x1x1xf32> into tensor<1x4x1x512x1x1xf32>
+  %1 = tensor.collapse_shape %0 [[0, 1, 2], [3], [4], [5]]
+    : tensor<1x4x1x512x1x1xf32> into tensor<4x512x1x1xf32>
+  return %1 : tensor<4x512x1x1xf32>
+}
+//       CHECK: func @fold_reshape_unit_dims
+//       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2], [3]]
+//  CHECK-SAME:   tensor<2048x1x1xf32> into tensor<4x512x1x1xf32>
+
+// -----
+
+func @expand_reshape_unit_dims(%arg0 : tensor<2048x1x2048xf32>)
+    -> tensor<4x512x1x512x4xf32> {
+  %0 = tensor.expand_shape %arg0 [[0, 1, 2, 3, 4], [5], [6, 7, 8]]
+    : tensor<2048x1x2048xf32> into tensor<1x4x1x512x1x1x512x1x4xf32>
+  %1 = tensor.collapse_shape %0 [[0, 1, 2], [3, 4], [5], [6, 7], [8]]
+    : tensor<1x4x1x512x1x1x512x1x4xf32> into tensor<4x512x1x512x4xf32>
+  return %1 : tensor<4x512x1x512x4xf32>
+}
+//       CHECK: func @expand_reshape_unit_dims
+//       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2], [3, 4]]
+//  CHECK-SAME:   tensor<2048x1x2048xf32> into tensor<4x512x1x512x4xf32>
+
+// -----
+
+func @fold_reshape_trailing_unit_dims(%arg0: tensor<2xf32>) -> tensor<2x1xf32> {
+  %0 = tensor.expand_shape %arg0 [[0, 1, 2]]
+      : tensor<2xf32> into tensor<2x1x1xf32>
+  %1 = tensor.collapse_shape %0 [[0], [1, 2]]
+      : tensor<2x1x1xf32> into tensor<2x1xf32>
+  return %1 : tensor<2x1xf32>
+}
+//       CHECK: func @fold_reshape_trailing_unit_dims
+//       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1]]
+//  CHECK-SAME:   tensor<2xf32> into tensor<2x1xf32>
+
+// -----
+
+func @collapse_reshape_unit_dims_dynamic(%arg0 : tensor<?x1x?x1x1x?x?x1x1xf32>)
+    -> tensor<?x?x?x?xf32> {
+  %0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3], [4], [5], [6, 7, 8]]
+    : tensor<?x1x?x1x1x?x?x1x1xf32> into tensor<?x?x1x1x?x?xf32>
+  %1 = tensor.collapse_shape %0 [[0], [1], [2, 3, 4], [5]]
+    : tensor<?x?x1x1x?x?xf32> into tensor<?x?x?x?xf32>
+  return %1 : tensor<?x?x?x?xf32>
+}
+//       CHECK: func @collapse_reshape_unit_dims_dynamic
+//       CHECK: tensor.collapse_shape
+//  CHECK-SAME:   [0], [1, 2], [3, 4, 5], [6, 7, 8]
+//  CHECK-SAME:   tensor<?x1x?x1x1x?x?x1x1xf32> into tensor<?x?x?x?xf32>
+
+// -----
+
+func @fold_reshape_trailing_unit_dims(%arg0: tensor<2xf32>) -> tensor<2x1xf32>
+{
+  %0 = tensor.expand_shape %arg0 [[0, 1, 2]]
+      : tensor<2xf32> into tensor<2x1x1xf32>
+  %1 = tensor.collapse_shape %0 [[0], [1, 2]]
+      : tensor<2x1x1xf32> into tensor<2x1xf32>
+  return %1 : tensor<2x1xf32>
+}
+//       CHECK: func @fold_reshape_trailing_unit_dims
+//       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1]]
+//  CHECK-SAME:   tensor<2xf32> into tensor<2x1xf32>
+
+// -----
+
+func @fold_reshape_trailing_unit_dims_dynamic(%arg0: tensor<1x1x?x1x1x1xf32>)
+    -> tensor<?xf32> {
+  %0 = tensor.collapse_shape %arg0 [[0, 1, 2], [3], [4], [5]]
+      : tensor<1x1x?x1x1x1xf32> into tensor<?x1x1x1xf32>
+  %1 = tensor.collapse_shape %0 [[0, 1, 2, 3]]
+      : tensor<?x1x1x1xf32> into tensor<?xf32>
+  return %1 : tensor<?xf32>
+}
+//       CHECK: func @fold_reshape_trailing_unit_dims_dynamic
+//       CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1, 2, 3, 4, 5]]
+//  CHECK-SAME:   tensor<1x1x?x1x1x1xf32> into tensor<?xf32>
+
+// -----
+
+func @fold_reshape_trailing_unit_dims(%arg0: tensor<12x42x1x1xf32>)
+    -> tensor<12x42xf32> {
+  %0 = tensor.expand_shape %arg0 [[0], [1], [2], [3, 4]]
+      : tensor<12x42x1x1xf32> into tensor<12x42x1x1x1xf32>
+  %1 = tensor.collapse_shape %0 [[0], [1, 2, 3, 4]]
+      : tensor<12x42x1x1x1xf32> into tensor<12x42xf32>
+  return %1 : tensor<12x42xf32>
+}
+//       CHECK: func @fold_reshape_trailing_unit_dims
+//       CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0], [1, 2, 3]]
+//  CHECK-SAME:   tensor<12x42x1x1xf32> into tensor<12x42xf32>
+
+// -----
+
+func @no_fold_reshapes(%arg0 : tensor<?x?x?xf32>) -> tensor<?x?xf32> {
+  %0 = tensor.expand_shape %arg0 [[0], [1], [2, 3]]
+      : tensor<?x?x?xf32> into tensor<?x?x1x?xf32>
+  %1 = tensor.collapse_shape %0 [[0], [1, 2, 3]]
+      : tensor<?x?x1x?xf32> into tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+// CHECK-LABEL: func @no_fold_reshapes
+//       CHECK:   tensor.expand_shape
+//       CHECK:   tensor.collapse_shape
+
+// -----
+
+func @no_fold_reshape_incompatible(%arg0 : tensor<4x6x8xf32>)
+    -> tensor<2x6x16xf32> {
+  %0 = tensor.expand_shape %arg0 [[0, 1], [2, 3], [4]]
+      : tensor<4x6x8xf32> into tensor<2x2x3x2x8xf32>
+  %1 = tensor.collapse_shape %0 [[0], [1, 2], [3, 4]]
+      : tensor<2x2x3x2x8xf32> into tensor<2x6x16xf32>
+  return %1 : tensor<2x6x16xf32>
+}
+// CHECK-LABEL: func @no_fold_reshape_incompatible
+//       CHECK:   tensor.expand_shape
+//       CHECK:   tensor.collapse_shape
+
+// -----
+
+func @no_fold_reshape_empty_expr(%arg0: tensor<3x2x2xf32>) -> tensor<12x1xf32> {
+  %0 = tensor.expand_shape %arg0 [[0], [1], [2, 3]]
+      : tensor<3x2x2xf32> into tensor<3x2x2x1xf32>
+  %1 = tensor.collapse_shape %0 [[0, 1, 2], [3]]
+      : tensor<3x2x2x1xf32> into tensor<12x1xf32>
+  return %1 : tensor<12x1xf32>
+}
+//      CHECK: func @no_fold_reshape_empty_expr
+// CHECK-SAME:    %[[ARG0:.+]]: tensor<3x2x2xf32>
+//      CHECK:    %[[RARG0:.+]] = tensor.expand_shape %[[ARG0]]
+// CHECK-SAME:      [0], [1], [2, 3]
+//      CHECK:    %[[RES:.+]] = tensor.collapse_shape %[[RARG0]]
+// CHECK-SAME:      [0, 1, 2], [3]
+//      CHECK:    return %[[RES:.+]] : tensor<12x1xf32>
+
+// -----
+
+func @reshape_splat_constant_int32() -> tensor<2x4x2xi32> {
+  %c0 = arith.constant dense<42> : tensor<2x8xi32>
+  %0 = tensor.expand_shape %c0 [[0], [1, 2]]
+      : tensor<2x8xi32> into tensor<2x4x2xi32>
+  return %0 : tensor<2x4x2xi32>
+}
+// CHECK-LABEL: @reshape_splat_constant_int32
+//       CHECK:   %[[CST:.*]] = arith.constant dense<{{.*}}> : tensor<2x4x2xi32>
+//   CHECK-NOT:   tensor.expand_shape
+//       CHECK:   return %[[CST]]
+
+// -----
+
+func @reshape_splat_constant_int16() -> tensor<2x4x2xi16> {
+  %c0 = arith.constant dense<42> : tensor<2x8xi16>
+  %0 = tensor.expand_shape %c0 [[0], [1, 2]]
+      : tensor<2x8xi16> into tensor<2x4x2xi16>
+  return %0 : tensor<2x4x2xi16>
+}
+// CHECK-LABEL: @reshape_splat_constant_int16
+//       CHECK:   %[[CST:.*]] = arith.constant dense<{{.*}}> : tensor<2x4x2xi16>
+//   CHECK-NOT:   tensor.expand_shape
+//       CHECK:   return %[[CST]]
+
+// -----
+
+func @reshape_splat_constant_float32() -> tensor<2x4x2xf32> {
+  %c0 = arith.constant dense<42.0> : tensor<2x8xf32>
+  %0 = tensor.expand_shape %c0 [[0], [1, 2]]
+      : tensor<2x8xf32> into tensor<2x4x2xf32>
+  return %0 : tensor<2x4x2xf32>
+}
+// CHECK-LABEL: @reshape_splat_constant_float32
+//       CHECK:   %[[CST:.*]] = arith.constant dense<{{.*}}> : tensor<2x4x2xf32>
+//   CHECK-NOT:   tensor.expand_shape
+//       CHECK:   return %[[CST]]
+
+// -----
+
+func @reshape_splat_constant_float64() -> tensor<2x4x2xf64> {
+  %c0 = arith.constant dense<42.0> : tensor<2x8xf64>
+  %0 = tensor.expand_shape %c0 [[0], [1, 2]]
+      : tensor<2x8xf64> into tensor<2x4x2xf64>
+  return %0 : tensor<2x4x2xf64>
+}
+// CHECK-LABEL: @reshape_splat_constant_float64
+//       CHECK:   %[[CST:.*]] = arith.constant dense<{{.*}}> : tensor<2x4x2xf64>
+//   CHECK-NOT:   tensor.expand_shape
+//       CHECK:   return %[[CST]]

mlir/test/Dialect/Tensor/invalid.mlir

 
 func @insert_slice_wrong_dynamic_type(%t1: tensor<?x4x4xf32>, %t2: tensor<8x16x4xf32>, %idx : index) {
   // expected-error @+1 {{expected type to be 'tensor<4x4x4xf32>' or a rank-reduced version. (size mismatch)}}
   %0 = tensor.insert_slice %t1 into %t2[0, 2, 0][4, 4, 4][1, 1, 1]
     : tensor<?x4x4xf32> into tensor<8x16x4xf32>
 
   return
 }
+
+// -----
+
+func @illegal_expanding_reshape_dynamic_tensor
+  (%arg0: tensor<?x?x?xf32>) -> tensor<?x?x?x4x?xf32> {
+  // expected-error @+1 {{invalid to have a single dimension (2) expanded into multiple dynamic dims (2,4)}}
+  %0 = tensor.expand_shape %arg0 [[0], [1], [2, 3, 4]]
+      : tensor<?x?x?xf32> into tensor<?x?x?x4x?xf32>
+  return %0 : tensor<?x?x?x4x?xf32>
+}
+
+// -----
+
+
+func @illegal_expanding_reshape_static_tensor
+    (%arg0: tensor<2x3x20xf32>) -> tensor<2x3x2x4x5xf32> {
+  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
+  %0 = tensor.expand_shape %arg0 [[0], [1], [2, 3, 4]]
+      : tensor<2x3x20xf32> into tensor<2x3x2x4x5xf32>
+  return %0 : tensor<2x3x2x4x5xf32>
+}
+
+// -----
+
+func @illegal_collapsing_reshape_static_tensor
+    (%arg0: tensor<2x3x2x4x5xf32>) -> tensor<2x3x20xf32> {
+  // expected-error @+1 {{expected dimension 2 of collapsed type to be static value of 40}}
+  %0 = tensor.collapse_shape %arg0 [[0], [1], [2, 3, 4]]
+      : tensor<2x3x2x4x5xf32> into tensor<2x3x20xf32>
+  return %0 : tensor<2x3x20xf32>
+}
+
+// -----
+
+func @illegal_expanding_reshape_mixed_tensor(%arg0 : tensor<?x?xf32>)
+    -> tensor<?x4x5xf32> {
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
+  %0 = tensor.expand_shape %arg0 [[0, 1], [2]]
+      : tensor<?x?xf32> into tensor<?x4x5xf32>
+  return %0 : tensor<?x4x5xf32>
+}
+
+// -----
+
+func @illegal_expanding_reshape_mixed_tensor_2(%arg0 : tensor<?x?xf32>)
+    -> tensor<?x4x5xf32> {
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
+  %0 = tensor.expand_shape %arg0 [[0], [1, 2]]
+      : tensor<?x?xf32> into tensor<?x4x5xf32>
+  return %0 : tensor<?x4x5xf32>
+}
+
+// -----
+
+func @illegal_collapsing_reshape_mixed_tensor(%arg0 : tensor<?x4x5xf32>) -> tensor<?x?xf32> {
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 5}}
+  %0 = tensor.collapse_shape %arg0 [[0, 1], [2]]
+      : tensor<?x4x5xf32> into tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// -----
+
+func @illegal_collapsing_reshape_mixed_tensor_2(%arg0 : tensor<?x4x5xf32>)
+    -> tensor<?x?xf32> {
+  // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
+  %0 = tensor.collapse_shape %arg0 [[0], [1, 2]]
+      : tensor<?x4x5xf32> into tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}

mlir/test/Dialect/Tensor/ops.mlir

func @slice(%t: tensor<8x16x4xf32>, %idx : index) {
   // CHECK: tensor.extract_slice
   // CHECK-SAME: tensor<8x16x4xf32> to tensor<4x4xf32>
   %3 = tensor.extract_slice %t[0, 2, 0][4, 1, 4][1, 1, 1]
     : tensor<8x16x4xf32> to tensor<4x4xf32>
 
   return
 }
 
+// -----
+
 // CHECK-LABEL: func @insert_slice({{.*}}) {
 func @insert_slice(
     %t: tensor<8x16x4xf32>,
     %td: tensor<8x?x4xf32>,
     %t2: tensor<16x32x8xf32>,
     %t3: tensor<4x4xf32>,
     %idx : index,
     %sz : index) {
func @insert_slice(
 
   // CHECK: tensor.insert_slice
   // CHECK-SAME: tensor<8x?x4xf32> into tensor<8x16x4xf32>
   %4 = tensor.insert_slice %td into %t[0, %idx, 0][8, %sz, 4][1, 1, 1]
     : tensor<8x?x4xf32> into tensor<8x16x4xf32>
 
   return
 }
+
+// -----
+
+func @tensor_reshape_zero_dim(%arg0 : tensor<1x1xf32>, %arg1 : tensor<f32>)
+    -> (tensor<f32>, tensor<1x1xf32>) {
+  %0 = tensor.collapse_shape %arg0 [] : tensor<1x1xf32> into tensor<f32>
+  %1 = tensor.expand_shape %0 [] : tensor<f32> into tensor<1x1xf32>
+  return %0, %1 : tensor<f32>, tensor<1x1xf32>
+}
+// CHECK-LABEL: func @tensor_reshape_zero_dim
+//       CHECK:   tensor.collapse_shape %{{.*}} [] : tensor<1x1xf32> into tensor<f32>
+//       CHECK:   tensor.expand_shape %{{.*}} [] : tensor<f32> into tensor<1x1xf32>
+
+func @legal_collapsing_reshape_dynamic_tensor
+  (%arg0: tensor<?x?x?x4x?xf32>) -> tensor<?x?x?xf32>
+{
+  %0 = tensor.collapse_shape %arg0 [[0], [1], [2, 3, 4]] :
+    tensor<?x?x?x4x?xf32> into tensor<?x?x?xf32>
+  return %0 : tensor<?x?x?xf32>
+}
+//      CHECK: func @legal_collapsing_reshape_dynamic_tensor
+//      CHECK:   tensor.collapse_shape
+// CHECK-SAME:    [0], [1], [2, 3, 4]

mlir/test/Integration/Dialect/Linalg/CPU/test-collapse-tensor.mlir

func @main() {
   // CHECK-NEXT: [-10.7136,    0.8428],
   // CHECK-NEXT: [16.4233,    9.4589]]]
   return
 }
 
 func private @print_memref_f32(%ptr : tensor<*xf32>)
 
 func @collapse_dynamic_shape(%arg0 : tensor<2x?x?x?xf32>) -> tensor<2x?x?xf32> {
-  %0 = linalg.tensor_collapse_shape %arg0 [[0], [1, 2], [3]]: tensor<2x?x?x?xf32> into tensor<2x?x?xf32>
+  %0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3]]: tensor<2x?x?x?xf32> into tensor<2x?x?xf32>
   return %0 : tensor<2x?x?xf32>
 }

mlir/test/Integration/Dialect/Linalg/CPU/test-expand-tensor.mlir

func @main() {
   // CHECK-NEXT{LITERAL}:   [[-10.7136,      0.8428]],
   // CHECK-NEXT{LITERAL}:   [[16.4233,      9.4589]]]]]
   return
 }
 
 func private @print_memref_f32(%ptr : tensor<*xf32>)
 
 func @expand_dynamic_shape(%arg0 : tensor<2x?x?xf32>) -> tensor<2x2x?x1x?xf32> {
-  %0 = linalg.tensor_expand_shape %arg0 [[0], [1, 2, 3], [4]]: tensor<2x?x?xf32> into tensor<2x2x?x1x?xf32>
+  %0 = tensor.expand_shape %arg0 [[0], [1, 2, 3], [4]]: tensor<2x?x?xf32> into tensor<2x2x?x1x?xf32>
   return %0 : tensor<2x2x?x1x?xf32>
 }

mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp

struct TestLinalgControlFuseByExpansion
   void runOnFunction() override {
     MLIRContext *context = &this->getContext();
     FuncOp funcOp = this->getFunction();
     RewritePatternSet fusionPatterns(context);
 
     linalg::ControlElementwiseOpsFusionFn controlReshapeFusionFn =
         [](const OpResult &producer, OpOperand &consumer) {
           if (auto collapseOp =
-                  producer.getDefiningOp<linalg::TensorCollapseShapeOp>()) {
+                  producer.getDefiningOp<tensor::CollapseShapeOp>()) {
             if (!collapseOp.src().getDefiningOp<linalg::LinalgOp>()) {
               return false;
             }
           }
           if (auto expandOp =
-                  dyn_cast<linalg::TensorExpandShapeOp>(consumer.getOwner())) {
+                  dyn_cast<tensor::ExpandShapeOp>(consumer.getOwner())) {
             if (expandOp->hasOneUse()) {
               OpOperand &use = *expandOp->getUses().begin();
               auto linalgOp = dyn_cast<linalg::LinalgOp>(use.getOwner());
               if (linalgOp && linalgOp.isOutputTensor(&use))
                 return true;
             }
           }
           return linalg::skipUnitDimReshape(producer, consumer);

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

gentbl_cc_library(
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Tensor/IR/TensorOps.td",
     deps = [":TensorOpsTdFiles"],
 )
 
 cc_library(
     name = "TensorDialect",
-    srcs = glob(
-        [
-            "lib/Dialect/Tensor/IR/*.cpp",
-            "lib/Dialect/Tensor/IR/*.h",
-        ],
-    ) + ["include/mlir/Transforms/InliningUtils.h"],
+    srcs = [
+        "include/mlir/Transforms/InliningUtils.h",
+        "lib/Dialect/Tensor/IR/TensorDialect.cpp",
+        "lib/Dialect/Tensor/IR/TensorOps.cpp",
     hdrs = ["include/mlir/Dialect/Tensor/IR/Tensor.h"],
     includes = ["include"],
     deps = [
         ":ArithmeticDialect",
         ":CastOpInterfaces",
         ":ControlFlowInterfaces",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":SideEffectInterfaces",
         ":StandardOps",
         ":Support",
         ":TensorOpsIncGen",
         ":ViewLikeInterface",
         "//llvm:Support",
     ],
 )
 
+cc_library(
+    name = "TensorInferTypeOpInterfaceImpl",
+    srcs = ["lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp"],
+    hdrs = ["include/mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"],
+    includes = ["include"],
+    deps = [
+        ":Affine",
+        ":IR",
+        ":InferTypeOpInterface",
+        ":StandardOps",
+        ":TensorDialect",
+        "//llvm:Support",
+    ],
+)
+
 gentbl_cc_library(
     name = "TensorPassIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
             [
                 "-gen-pass-decls",
                 "-name=Tensor",
deps = [
         ":SparseTensor",
         ":SparseTensorTransforms",
         ":StandardOps",
         ":StandardOpsTransforms",
         ":StandardOpsTransformsPassIncGen",
         ":StandardToLLVM",
         ":StandardToSPIRV",
         ":TensorDialect",
+        ":TensorInferTypeOpInterfaceImpl",
         ":TensorTransforms",
         ":TosaDialect",
         ":TosaToLinalg",
         ":Transforms",
         ":TransformsPassIncGen",
         ":VectorOps",
         ":VectorToLLVM",
         ":VectorToROCDL",