[mlir] Remove unnecessary canonicalization from Linalg Detensorize.cpp

[pifon2a]

Summary

After https://reviews.llvm.org/D115821 it became possible to create
tensor<elem_type> with a single tensor.from_elements operation
without collapsing tensor shape from tensor<1xelem_type> to
tensor<elem_type>.

[nicolasvasilache]

Nice!

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

 #include <utility>
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 static Value sourceMaterializationCallback(OpBuilder &builder, Type type,
                                            ValueRange inputs, Location loc) {
   assert(inputs.size() == 1);
-  if (inputs[0].getType().isa<TensorType>())
+  auto inputType = inputs[0].getType();
+  if (inputType.isa<TensorType>())
     return nullptr;
 
   // A detensored value is converted back by creating a new tensor from its
   // element(s).
-  auto createNewTensorOp =
-      builder.create<tensor::FromElementsOp>(loc, inputs[0]);
-
-  // FromElementsOp results in a tensor<1xdtype>, we need to reshape that to
-  // a tensor<dtype> instead.
-  return builder.create<tensor::CollapseShapeOp>(
-      loc, type, createNewTensorOp, ArrayRef<ReassociationExprs>{});
+  return builder.create<tensor::FromElementsOp>(
+      loc, RankedTensorType::get({}, inputType), inputs[0]);
 }
 
 namespace {
 /// Defines the criteria a TensorType must follow in order to be considered
 /// "detensorable".
 ///
 /// NOTE: For now, only 0-D tensors are supported.
 ///
addTargetMaterialization([](OpBuilder &builder, Type type,
       return builder.create<tensor::ExtractOp>(loc, inputs[0], ValueRange{});
     });
 
     addSourceMaterialization(sourceMaterializationCallback);
     addArgumentMaterialization(sourceMaterializationCallback);
   }
 };
 
-/// Canonicalizes the pattern of the form
-///
-/// %tensor = tensor.from_elements(%element) : (i32) -> tensor<1xi32>
-/// %reshaped_tensor = tensor.collapse_shape %tensor []
-///     : tensor<1xi32> into tensor<i32>
-/// %extracted_element = tensor.extract %reshaped_tensor[] : tensor<i32>
-///
-/// to just %element.
-struct ExtractFromReshapeFromElements
-    : public OpRewritePattern<tensor::ExtractOp> {
-  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<tensor::CollapseShapeOp>();
-    if (tensorReshape == nullptr)
-      return failure();
-
-    auto tensorFromElements =
-        tensorReshape.getOperand()
-            .getDefiningOp<mlir::tensor::FromElementsOp>();
-    if (tensorFromElements == nullptr)
-      return failure();
-
-    rewriter.replaceOp(extract, tensorFromElements.getOperand(0));
-    return success();
-  }
-};
-
 /// @see LinalgDetensorize in Linalg/Passes.td for more details.
 struct LinalgDetensorize : public LinalgDetensorizeBase<LinalgDetensorize> {
   LinalgDetensorize() = default;
 
   class CostModel {
   public:
     virtual ~CostModel() = default;
 
void runOnOperation() override {
     populateBranchOpInterfaceTypeConversionPattern(patterns, typeConverter,
                                                    shouldConvertBranchOperand);
 
     if (failed(
             applyFullConversion(getOperation(), target, std::move(patterns))))
       signalPassFailure();
 
     RewritePatternSet canonPatterns(context);
-    canonPatterns.add<ExtractFromReshapeFromElements>(context);
+    tensor::FromElementsOp::getCanonicalizationPatterns(canonPatterns, context);
     if (failed(applyPatternsAndFoldGreedily(getOperation(),
                                             std::move(canonPatterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createLinalgDetensorizePass() {
   return std::make_unique<LinalgDetensorize>();
 }

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:.*]] = tensor.collapse_shape %[[new_tensor_res]]
-// CHECK:         return %[[reshaped_tensor_res]]
+// CHECK:         return %[[new_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
     %2 = arith.addf %arg3, %arg4 : f32

 // 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:.*]] = tensor.collapse_shape %[[new_tensor_res]]
-// CHECK:         return %[[reshaped_tensor_res]]
+// CHECK:         return %[[new_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:.*]] = tensor.collapse_shape %[[new_tensor_res]]
-// CHECK:         return %[[reshaped_tensor_res]]
+// CHECK:         return %[[new_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:.*]] = tensor.collapse_shape %[[new_tensor_res]]
-// CHECK:         return %[[reshaped_tensor_res]]
+// CHECK:         return %[[new_tensor_res]]

mlir/test/Dialect/Linalg/detensorize_br_operands.mlir

 // RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -pass-pipeline="builtin.func(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 = tensor.collapse_shape %arg0_t [] : tensor<1xi1> into tensor<i1>
-
-  %arg1_t = tensor.from_elements %arg1 : tensor<1xi32>
-  %arg1_t2 = tensor.collapse_shape %arg1_t [] : tensor<1xi32> into tensor<i32>
+  %arg0_t = tensor.from_elements %arg0 : tensor<i1>
+  %arg1_t = tensor.from_elements %arg1 : 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>)
+    ins(%arg0_t : tensor<i1>)
     outs(%2 : tensor<i8>) {
   ^bb0(%arg2: i1, %arg3: i8):  // no predecessors
     %10 = arith.extui %arg2 : i1 to i8
     linalg.yield %10 : i8
   } -> tensor<i8>
   %4 = tensor.extract %3[] : tensor<i8>
   %5 = arith.trunci %4 : i8 to i1
-  cond_br %5, ^bb1, ^bb2(%arg1_t2 : tensor<i32>)
+  cond_br %5, ^bb1, ^bb2(%arg1_t : tensor<i32>)
 ^bb1:
   %6 = linalg.init_tensor [] : tensor<i32>
   %7 = linalg.generic
     {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []}
-    ins(%arg1_t2, %cst : tensor<i32>, tensor<i32>)
+    ins(%arg1_t, %cst : tensor<i32>, tensor<i32>)
     outs(%6 : tensor<i32>) {
   ^bb0(%arg2: i32, %arg3: i32, %arg4: i32):  // no predecessors
     %10 = arith.addi %arg2, %arg3 : i32
     linalg.yield %10 : i32
   } -> tensor<i32>
   br ^bb2(%7 : tensor<i32>)
 ^bb2(%8: tensor<i32>):
   return %8 : tensor<i32>
 }
 
 // CHECK-LABEL:  func @if_true_test
 // 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:.*]] = tensor.collapse_shape
+// CHECK-NEXT:     %[[func_res:.*]] = tensor.from_elements
 // CHECK-NEXT:     return %[[func_res]]

mlir/test/Dialect/Linalg/detensorize_if.mlir

 // RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -pass-pipeline="builtin.func(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 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
+  %0 = tensor.from_elements %c0 : tensor<i32>
   %c10 = arith.constant 10 : i32
-  %1 = tensor.from_elements %c10 : tensor<1xi32>
-  %reshaped1 = tensor.collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
-  br ^bb1(%reshaped0 : tensor<i32>)
+  %1 = tensor.from_elements %c10 : tensor<i32>
+  br ^bb1(%0 : 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>)
+    ins(%2, %1 : 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>
   cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)
 

 // CHECK-NEXT:     br ^[[bb1:.*]](%{{.*}}: i32)
 // 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:     tensor.collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
+// CHECK-NEXT:     tensor.from_elements %{{.*}} : 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 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
+  %0 = tensor.from_elements %c0 : tensor<i32>
   %c10 = arith.constant 10 : i32
-  %1 = tensor.from_elements %c10 : tensor<1xi32>
-  %reshaped1 = tensor.collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
-  br ^bb1(%reshaped0 : tensor<i32>)
+  %1 = tensor.from_elements %c10 : tensor<i32>
+  br ^bb1(%0 : 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>)
+    ins(%2, %1 : 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>
   cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<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:     br ^[[bb4:.*]](%{{.*}} : i32)
 // CHECK-NEXT:   ^[[bb4]](%{{.*}}: i32)
-// CHECK-NEXT:     tensor.from_elements %{{.*}} : tensor<1xi32>
-// CHECK-NEXT:     tensor.collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
+// CHECK-NEXT:     tensor.from_elements %{{.*}} : 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 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
+  %0 = tensor.from_elements %c0 : tensor<i32>
   %c10 = arith.constant 10 : i32
-  %1 = tensor.from_elements %c10 : tensor<1xi32>
-  %reshaped1 = tensor.collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
-  br ^bb1(%reshaped0 : tensor<i32>)
+  %1 = tensor.from_elements %c10 : tensor<i32>
+  br ^bb1(%0 : 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>)
+    ins(%2, %1 : 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 = tensor.collapse_shape %12 [] : tensor<1xi32> into tensor<i32>
+  %12 = tensor.from_elements %c10 : tensor<i32>
   %7 = linalg.init_tensor [] : tensor<i32>
   %8 = linalg.generic #attrs
-    ins(%6, %reshaped12 : tensor<i32>, tensor<i32>)
+    ins(%6, %12 : 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>
   br ^bb3(%8 : tensor<i32>)
 
 ^bb3(%10: tensor<i32>):  // pred: ^bb1
   return %10 : tensor<i32>
 }
 
 // CHECK-LABEL:  func @main()
 // CHECK-NEXT:     arith.constant 0
 // CHECK-NEXT:     arith.constant 10
 // CHECK-NEXT:     br ^[[bb1:.*]](%{{.*}}: 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:     tensor.collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
+// CHECK-NEXT:     tensor.from_elements %{{.*}} : tensor<i32>
 // CHECK-NEXT:     return %{{.*}}
 // CHECK-NEXT:   }

mlir/test/Dialect/Linalg/detensorize_trivial.mlir

 // RUN: mlir-opt %s -pass-pipeline="builtin.func(linalg-detensorize{aggressive-mode})" | FileCheck %s -check-prefix=DET-ALL
 // RUN: mlir-opt %s -pass-pipeline="builtin.func(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 = tensor.collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
+  %1 = tensor.from_elements %c10 : tensor<i32>
   %3 = linalg.init_tensor [] : tensor<i1>
   %4 = linalg.generic #attrs
-    ins(%farg0, %reshaped1 : tensor<i32>, tensor<i32>)
+    ins(%farg0, %1 : 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:    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:         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:         tensor.collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
+// DET-CF:         tensor.from_elements %{{.*}} : tensor<i32>
 // DET-CF:         return %{{.*}} : tensor<i32>

mlir/test/Dialect/Linalg/detensorize_while_impure_cf.mlir

 // DET-ALL:         ^bb0(%{{.*}}: i32, %{{.*}}: i32):  // no predecessors
 // 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:         tensor.collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
+// DET-ALL:         tensor.from_elements %{{.*}} : 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:         tensor.collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
+// DET-ALL:         tensor.from_elements %{{.*}} : 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>) {