[MLIR][Shape] Simplify shape lowering

[frgossen]

Summary

Depends On D84160

[pifon2a]

mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp

203

nit: s/auto/Value here and where possible.

[frgossen]

Address comment

mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp

ShapeOfOpConverter::matchAndRewrite(ShapeOfOp op, ArrayRef<Value> operands,
   if (tensorTy.isa<RankedTensorType>())
     return failure();
 
   // Allocate stack memory.
   auto loc = op.getLoc();
   auto rankVal = rewriter.create<mlir::RankOp>(loc, tensorVal);
   auto i64Ty = rewriter.getI64Type();
   auto memTy = MemRefType::get({ShapedType::kDynamicSize}, i64Ty);
-  auto memVal = rewriter.create<AllocaOp>(loc, memTy, ValueRange({rankVal}));
+  auto memVal = rewriter.create<AllocaOp>(loc, memTy, ValueRange{rankVal});
 
   // Copy shape extents to stack-allocated memory.
   auto zeroVal = rewriter.create<ConstantIndexOp>(loc, 0);
   auto oneVal = rewriter.create<ConstantIndexOp>(loc, 1);
   rewriter.create<scf::ForOp>(
       loc, zeroVal, rankVal, oneVal, llvm::None,
       [&](OpBuilder &b, Location loc, Value iVal, ValueRange args) {
         auto dimVal = rewriter.create<DimOp>(loc, tensorVal, iVal);
         auto dimIntVal = rewriter.create<IndexCastOp>(loc, dimVal, i64Ty);
         rewriter.create<StoreOp>(loc, dimIntVal, memVal, ValueRange{iVal});
         rewriter.create<scf::YieldOp>(loc);
       });
 
   // Load extents to tensor value.
-  auto shapeIntVal = rewriter.create<TensorLoadOp>(loc, memVal);
-  auto indexTy = rewriter.getIndexType();
-  auto shapeTy = RankedTensorType::get({ShapedType::kDynamicSize}, indexTy);
-  rewriter.replaceOpWithNewOp<IndexCastOp>(op.getOperation(), shapeIntVal,
-                                           shapeTy);
+  auto extentTensorInt = rewriter.create<TensorLoadOp>(loc, memVal);
+  rewriter.replaceOpWithNewOp<IndexCastOp>(op.getOperation(), extentTensorInt,
+                                           op.getType());
   return success();
 }
 
 namespace {
 struct ConvertShapeToSCFPass
     : public ConvertShapeToSCFBase<ConvertShapeToSCFPass> {
   void runOnFunction() override;
 };