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

Co-authored-by: Fabian Mora <[email protected]>

Author: nicolasvasilache

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

@@ -366,9 +366,8 @@ def ReifyRankedShapedTypeOpInterface :
 
         `reifiedReturnShapes` is populated with one vector per op result. Each
         of those vectors contains an OpFoldResult for each dimension of the
-        shaped type. In case a dimension in the type is static, the
-        corresponding entry is an IntegerAttr. Otherwise, it is a Value. The
-        given builder may be used to insert ops that compute result shapes.
+        shaped type. The given builder may be used to insert ops that compute
+        result shapes.
 
         If the shape of a particular result cannot be computed it must be empty.
       }],

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

@@ -7,13 +7,12 @@
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AsmState.h"
-#include "mlir/IR/BuiltinOps.h"
-#include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
@@ -195,9 +194,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(