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

Co-authored-by: Fabian Mora <fmora.dev@gmail.com>

Author: nicolasvasilache

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -195,9 +196,12 @@ FailureOr<Value> bufferization::allocateTensorForShapedValue(
         reifiedShapes = true;
         auto &shape =
             resultDims[llvm::cast<OpResult>(shapedValue).getResultNumber()];
-        for (const auto &dim : enumerate(tensorType.getShape()))
-          if (ShapedType::isDynamic(dim.value()))
-            dynamicSizes.push_back(cast<Value>(shape[dim.index()]));
+        for (const auto &dim : enumerate(tensorType.getShape())) {
+          if (ShapedType::isDynamic(dim.value())) {
+            dynamicSizes.push_back(
+                getValueOrCreateConstantIndexOp(b, loc, shape[dim.index()]));
+          }
+        }
       }
     }
 

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

@@ -3809,7 +3809,6 @@ PadOp::reifyResultShapes(OpBuilder &b,
     Value dim = b.createOrFold<tensor::DimOp>(
         loc, getSource(), b.create<arith::ConstantIndexOp>(loc, i));
 
-    affine::AffineBuilder ab(b, loc);
     AffineExpr d0, d1, d2;
     bindDims(b.getContext(), d0, d1, d2);
     reifiedReturnShapes[0][i] = affine::makeComposedFoldedAffineApply(