[mlir][NFC] Clean up builder usage around constants/non-foldable ops

[springerm]

Summary

• Use create instead of createOrFold for constant ops. Constants cannot be folded any further.
• Use create instead of createOrFold for ops that do not have a folder.
• Use C++ op builders that take an int instead of creating a ConstantIndexOp.
• Create tensor::DimOp instead of linalg::createOrFoldDimOp when it is certain that the operand is a tensor.

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

 } // namespace mlir
 
 using namespace mlir;
 using namespace mlir::amdgpu;
 
 static Value createI32Constant(ConversionPatternRewriter &rewriter,
                                Location loc, int32_t value) {
   Type llvmI32 = rewriter.getI32Type();
-  return rewriter.createOrFold<LLVM::ConstantOp>(loc, llvmI32, value);
+  return rewriter.create<LLVM::ConstantOp>(loc, llvmI32, value);
 }
 
 namespace {
 /// Define lowering patterns for raw buffer ops
 template <typename GpuOp, typename Intrinsic>
 struct RawBufferOpLowering : public ConvertOpToLLVMPattern<GpuOp> {
   RawBufferOpLowering(LLVMTypeConverter &converter, Chipset chipset)
       : ConvertOpToLLVMPattern<GpuOp>(converter), chipset(chipset) {}

 static Value mfmaConcatIfNeeded(ConversionPatternRewriter &rewriter,
                                 Location loc, Value input) {
   Type inputType = input.getType();
   if (auto vectorType = dyn_cast<VectorType>(inputType)) {
     if (!vectorType.getElementType().isInteger(8))
       return input;
     int64_t numBytes = vectorType.getNumElements();
     Type destType = rewriter.getIntegerType(numBytes * 8);
-    Value result = rewriter.createOrFold<LLVM::ConstantOp>(
+    Value result = rewriter.create<LLVM::ConstantOp>(
         loc, destType, rewriter.getIntegerAttr(destType, 0));
     for (int64_t i = 0; i < numBytes; ++i) {
       Value idxConst = createI32Constant(rewriter, loc, i);
       Value element =
           rewriter.create<LLVM::ExtractElementOp>(loc, input, idxConst);
       Value extended = rewriter.create<LLVM::ZExtOp>(loc, destType, element);
-      Value shiftConst = rewriter.createOrFold<LLVM::ConstantOp>(
+      Value shiftConst = rewriter.create<LLVM::ConstantOp>(
           loc, destType, rewriter.getIntegerAttr(destType, i * 8));
       Value shifted = rewriter.create<LLVM::ShlOp>(loc, extended, shiftConst);
       result = rewriter.create<LLVM::OrOp>(loc, result, shifted);
     }
     return result;
   }
   return input;
 }

mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp

matchAndRewrite(bufferization::CloneOp op, OpAdaptor adaptor,
 
     // Transform a clone operation into alloc + copy operation and pay
     // attention to the shape dimensions.
     Location loc = op->getLoc();
     SmallVector<Value, 4> dynamicOperands;
     for (int i = 0; i < memrefType.getRank(); ++i) {
       if (!memrefType.isDynamicDim(i))
         continue;
-      Value size = rewriter.createOrFold<arith::ConstantIndexOp>(loc, i);
-      Value dim =
-          rewriter.createOrFold<memref::DimOp>(loc, op.getInput(), size);
+      Value dim = rewriter.createOrFold<memref::DimOp>(loc, op.getInput(), i);
       dynamicOperands.push_back(dim);
     }
 
     // Allocate a memref with identity layout.
     Value alloc = rewriter.create<memref::AllocOp>(op->getLoc(), allocType,
                                                    dynamicOperands);
     // Cast the allocation to the specified type if needed.
     if (memrefType != allocType)

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

llvm::SmallVector<Value> genericOpOutputs = {
         createZeroTensor(rewriter, loc, outputType, dynamicSizes)};
 
     // Indexing maps for input and output tensors
     auto indexingMaps = AffineMap::inferFromExprList(llvm::ArrayRef{
         affineDimsExpr(rewriter, 0, 3, 4), affineDimsExpr(rewriter, 0, 1, 2),
         affineDimsExpr(rewriter, 0, 1, 2)});
 
     // Width and height dimensions of the original input.
-    auto dimH = linalg::createOrFoldDimOp(rewriter, loc, input, 1);
-    auto dimW = linalg::createOrFoldDimOp(rewriter, loc, input, 2);
+    auto dimH = rewriter.createOrFold<tensor::DimOp>(loc, input, 1);
+    auto dimW = rewriter.createOrFold<tensor::DimOp>(loc, input, 2);
 
     // Constants and dimension sizes
     auto twoPiAttr = rewriter.getFloatAttr(elementType, 6.283185307179586);
     auto twoPi = rewriter.create<arith::ConstantOp>(loc, twoPiAttr);
     auto constH = castIndexToFloat(rewriter, loc, elementType, dimH);
     auto constW = castIndexToFloat(rewriter, loc, elementType, dimW);
 
     auto buildBody = [&](OpBuilder &builder, Location loc, ValueRange args) {

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

if (memref::CastOp::areCastCompatible(srcType, destType) &&
     return casted;
   }
 
   auto loc = value.getLoc();
   SmallVector<Value, 4> dynamicOperands;
   for (int i = 0; i < destType.getRank(); ++i) {
     if (destType.getShape()[i] != ShapedType::kDynamic)
       continue;
-    auto index = b.createOrFold<arith::ConstantIndexOp>(loc, i);
-    Value size = b.create<memref::DimOp>(loc, value, index);
+    Value size = b.create<memref::DimOp>(loc, value, i);
     dynamicOperands.push_back(size);
   }
   // TODO: Use alloc/memcpy callback from BufferizationOptions if called via
   // BufferizableOpInterface impl of ToMemrefOp.
   Value copy = b.create<memref::AllocOp>(loc, destType, dynamicOperands);
   b.create<memref::CopyOp>(loc, value, copy);
   return copy;
 }

mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp

if (rank < GPUDialect::getNumWorkgroupDimensions()) {
     steps.resize(extraLoops, one);
   }
 
   // Add existing bounds.
   lbs.append(rank, zero);
   ubs.reserve(lbs.size());
   steps.reserve(lbs.size());
   for (auto idx = 0; idx < rank; ++idx) {
-    ubs.push_back(b.createOrFold<memref::DimOp>(
-        from, b.create<arith::ConstantIndexOp>(idx)));
+    ubs.push_back(b.createOrFold<memref::DimOp>(from, idx));
     steps.push_back(one);
   }
 
   // Obtain thread identifiers and block sizes, necessary to map to them.
   auto indexType = b.getIndexType();
   SmallVector<Value, 3> threadIds, blockDims;
   for (auto dim : {gpu::Dimension::x, gpu::Dimension::y, gpu::Dimension::z}) {
     threadIds.push_back(b.create<gpu::ThreadIdOp>(indexType, dim));

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

static Value allocBuffer(ImplicitLocOpBuilder &b,
   if (alignment.has_value())
     alignmentAttr = b.getI64IntegerAttr(alignment.value());
 
   // Static buffer.
   if (auto cst = allocSize.getDefiningOp<arith::ConstantIndexOp>()) {
     auto staticBufferType =
         MemRefType::get(width * cst.value(), b.getIntegerType(8));
     if (options.useAlloca) {
-      return b.createOrFold<memref::AllocaOp>(staticBufferType, ValueRange{},
-                                              alignmentAttr);
+      return b.create<memref::AllocaOp>(staticBufferType, ValueRange{},
     }
-    return b.createOrFold<memref::AllocOp>(staticBufferType, ValueRange{},
-                                           alignmentAttr);
+    return b.create<memref::AllocOp>(staticBufferType, ValueRange{},
   }
 
   // Fallback dynamic buffer.
   auto dynamicBufferType =
       MemRefType::get(ShapedType::kDynamic, b.getIntegerType(8));
   Value mul = b.createOrFold<arith::MulIOp>(
       b.create<arith::ConstantIndexOp>(width), allocSize);
   if (options.useAlloca)
     return b.create<memref::AllocaOp>(dynamicBufferType, mul, alignmentAttr);
   return b.create<memref::AllocOp>(dynamicBufferType, mul, alignmentAttr);
 }
 
 /// Default allocation callback function. This allocates a promoted buffer when
 /// no call back to do so is provided. The default is to allocate a
 /// memref<..xi8> and return a view to get a memref type of shape
 /// boundingSubViewSize.
 static std::optional<Value> defaultAllocBufferCallBack(
     const LinalgPromotionOptions &options, OpBuilder &builder,
     memref::SubViewOp subView, ArrayRef<Value> boundingSubViewSize,
     std::optional<unsigned> alignment, DataLayout &layout) {
   ShapedType viewType = subView.getType();
   ImplicitLocOpBuilder b(subView.getLoc(), builder);
-  auto zero = b.createOrFold<arith::ConstantIndexOp>(0);
-  auto one = b.createOrFold<arith::ConstantIndexOp>(1);
+  auto zero = b.create<arith::ConstantIndexOp>(0);
+  auto one = b.create<arith::ConstantIndexOp>(1);
 
   Value allocSize = one;
   for (const auto &size : llvm::enumerate(boundingSubViewSize))
     allocSize = b.createOrFold<arith::MulIOp>(allocSize, size.value());
   Value buffer = allocBuffer(b, options, viewType.getElementType(), allocSize,
                              layout, alignment);
   SmallVector<int64_t, 4> dynSizes(boundingSubViewSize.size(),
                                    ShapedType::kDynamic);

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

LogicalResult matchAndRewrite(OpTy dimOp,
     Value resultShape = reifiedResultShapes[dimValue.getResultNumber()];
     auto resultShapeType = dyn_cast<RankedTensorType>(resultShape.getType());
     if (!resultShapeType || !isa<IndexType>(resultShapeType.getElementType()))
       return failure();
 
     Location loc = dimOp->getLoc();
     rewriter.replaceOpWithNewOp<tensor::ExtractOp>(
         dimOp, resultShape,
-        rewriter.createOrFold<arith::ConstantIndexOp>(loc, *dimIndex));
+        rewriter.create<arith::ConstantIndexOp>(loc, *dimIndex).getResult());
     return success();
   }
 };
 
 /// Fold dim of an operation that implements the InferShapedTypeOpInterface
 template <typename OpTy>
 struct DimOfReifyRankedShapedTypeOpInterface : public OpRewritePattern<OpTy> {
   using OpRewritePattern<OpTy>::OpRewritePattern;