Changes in rocMLIR due to llvm/llvm-project#142986

Author: umangyadav

mlir/lib/Conversion/TosaToRock/TosaToRock.cpp

@@ -12,7 +12,9 @@
 
 #include "mlir/Conversion/TosaToRock/TosaToRock.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Rock/IR/Rock.h"
@@ -1195,8 +1197,8 @@ struct ConvElementwiseGemmRewritePattern
       RankedTensorType resTensorType = cast<RankedTensorType>(res.getType());
       MemRefType resMemRefType = MemRefType::get(
           resTensorType.getShape(), resTensorType.getElementType());
-      Value resMemref =
-          rewriter.create<bufferization::ToBufferOp>(loc, resMemRefType, res);
+      Value resMemref = rewriter.create<bufferization::ToBufferOp>(
+          loc, cast<mlir::bufferization::BufferLikeType>(resMemRefType), res);
       Value outMemref =
           preSecondGemmElemwiseBlock->addArgument(resMemRefType, loc);
       rewriter.create<memref::CopyOp>(loc, resMemref, outMemref);
@@ -1279,8 +1281,8 @@ struct GemmElementwiseGemmRewritePattern
       RankedTensorType resTensorType = cast<RankedTensorType>(res.getType());
       MemRefType resMemRefType = MemRefType::get(
           resTensorType.getShape(), resTensorType.getElementType());
-      Value resMemref =
-          rewriter.create<bufferization::ToBufferOp>(loc, resMemRefType, res);
+      Value resMemref = rewriter.create<bufferization::ToBufferOp>(
+          loc, cast<mlir::bufferization::BufferLikeType>(resMemRefType), res);
       Value outMemref =
           preSecondGemmElemwiseBlock->addArgument(resMemRefType, loc);
       rewriter.create<memref::CopyOp>(loc, resMemref, outMemref);
@@ -1714,7 +1716,7 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
       MemRefType resMemRefType = MemRefType::get(
           resTensorType.getShape(), resTensorType.getElementType());
       Value resMemref =
-          rewriter.create<bufferization::ToBufferOp>(loc, resMemRefType, res);
+          rewriter.create<bufferization::ToBufferOp>(loc, cast<bufferization::BufferLikeType>(resMemRefType), res);
       Value outMemref =
           preSoftmaxElemwiseBlock->addArgument(resMemRefType, loc);
       rewriter.create<memref::CopyOp>(loc, resMemref, outMemref);

mlir/lib/Dialect/Rock/utility/builderUtils.cpp

@@ -209,7 +209,8 @@ Value getAsTensor(OpBuilder &builder, Location loc, mlir::Value value,
                   bool isWritable) {
   constexpr bool isRestrict{true};
   Value origTensor = builder.create<bufferization::ToTensorOp>(
-      loc, value.getType(), value, isRestrict, isWritable);
+      loc, memref::getTensorTypeFromMemRefType(value.getType()), value,
+      isRestrict, isWritable);
   return origTensor;
 }
 

mlir/tools/rocmlir-gen/rocmlir-gen.cpp

@@ -15,6 +15,7 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
@@ -2911,8 +2912,9 @@ static func::FuncOp createGpuAttentionKernel(ModuleOp module,
     MemRefType resMemRefType =
         MemRefType::get({qShape[0], sequenceLengthQ, sequenceLengthK},
                         cast<ShapedType>(qkTensor.getType()).getElementType());
-    Value resMemref =
-        builder.create<bufferization::ToBufferOp>(loc, resMemRefType, qkTensor);
+    Value resMemref = builder.create<bufferization::ToBufferOp>(
+        loc, cast<mlir::bufferization::BufferLikeType>(resMemRefType),
+        qkTensor);
     Value outMemref = preSoftmaxElemwiseBlock->addArgument(resMemRefType, loc);
     builder.create<memref::CopyOp>(loc, resMemref, outMemref);
     builder.create<rock::YieldOp>(loc);
@@ -3002,8 +3004,9 @@ createGpuConvElementwiseGemmKernel(ModuleOp module, const GenParams &params) {
     MemRefType resMemRefType =
         MemRefType::get({aShape[0], firstGemmSize.m, firstGemmSize.n},
                         cast<ShapedType>(abTensor.getType()).getElementType());
-    Value resMemref =
-        builder.create<bufferization::ToBufferOp>(loc, resMemRefType, abTensor);
+    Value resMemref = builder.create<bufferization::ToBufferOp>(
+        loc, cast<mlir::bufferization::BufferLikeType>(resMemRefType),
+        abTensor);
     Value outMemref = preSecondGemmBlock->addArgument(resMemRefType, loc);
     builder.create<memref::CopyOp>(loc, resMemref, outMemref);
     builder.create<rock::YieldOp>(loc);
@@ -3098,8 +3101,9 @@ createGpuGemmElementwiseGemmKernel(ModuleOp module, const GenParams &params) {
     MemRefType resMemRefType =
         MemRefType::get({aShape[0], gemmM, gemmN},
                         cast<ShapedType>(abTensor.getType()).getElementType());
-    Value resMemref =
-        builder.create<bufferization::ToBufferOp>(loc, resMemRefType, abTensor);
+    Value resMemref = builder.create<bufferization::ToBufferOp>(
+        loc, cast<mlir::bufferization::BufferLikeType>(resMemRefType),
+        abTensor);
     Value outMemref = preSecondGemmBlock->addArgument(resMemRefType, loc);
     builder.create<memref::CopyOp>(loc, resMemref, outMemref);
     builder.create<rock::YieldOp>(loc);
@@ -3280,7 +3284,7 @@ createCpuConvElementwiseGemmKernelWithMlir(ModuleOp module,
                                           bool isWritable = false) {
     constexpr bool isRestrict{true};
     Value flatTensor = builder.create<bufferization::ToTensorOp>(
-        loc, block->getArgument(blockArgIndex).getType(),
+        loc, memref::getTensorTypeFromMemRefType(block->getArgument(blockArgIndex).getType()),
         block->getArgument(blockArgIndex), isRestrict, isWritable);
     ArrayRef<int64_t> origShape =
         cast<ShapedType>(argTypes[blockArgIndex]).getShape();
@@ -3418,11 +3422,11 @@ createCpuConvElementwiseGemmKernelWithMlir(ModuleOp module,
   }
 
   Value output = block->getArguments().back();
-  auto outputType = cast<MemRefType>(output.getType());
+  auto outputType = cast<bufferization::BufferLikeType>(output.getType());
 
   ImplicitLocOpBuilder implicitBuilder(loc, builder);
-  auto shapeValue =
-      tosa::getTosaConstShape(implicitBuilder, outputType.getShape());
+  auto shapeValue = tosa::getTosaConstShape(
+      implicitBuilder, cast<ShapedType>(outputType).getShape());
   auto flatResultTensor =
       builder.create<tosa::ReshapeOp>(loc, resultTensor, shapeValue);
 
@@ -3460,7 +3464,7 @@ createCpuGemmElementwiseGemmKernelWithMlir(ModuleOp module,
                                           bool isWritable = false) {
     constexpr bool isRestrict{true};
     Value flatTensor = builder.create<bufferization::ToTensorOp>(
-        loc, block->getArgument(blockArgIndex).getType(),
+        loc, memref::getTensorTypeFromMemRefType(block->getArgument(blockArgIndex).getType()),
         block->getArgument(blockArgIndex), isRestrict, isWritable);
     ArrayRef<int64_t> origShape =
         cast<ShapedType>(argTypes[blockArgIndex]).getShape();
@@ -3534,11 +3538,11 @@ createCpuGemmElementwiseGemmKernelWithMlir(ModuleOp module,
   }
 
   Value output = block->getArguments().back();
-  auto outputType = cast<MemRefType>(output.getType());
+  auto outputType = cast<mlir::bufferization::BufferLikeType>(output.getType());
 
   ImplicitLocOpBuilder implicitBuilder(loc, builder);
-  auto shapeValue =
-      tosa::getTosaConstShape(implicitBuilder, outputType.getShape());
+  auto shapeValue = tosa::getTosaConstShape(
+      implicitBuilder, cast<ShapedType>(outputType).getShape());
   auto flatResultTensor =
       builder.create<tosa::ReshapeOp>(loc, resultTensor, shapeValue);
 
@@ -3576,7 +3580,7 @@ static func::FuncOp createCpuAttentionKernelWithMlir(ModuleOp module,
                                           bool isWritable = false) {
     constexpr bool isRestrict{true};
     Value flatTensor = builder.create<bufferization::ToTensorOp>(
-        loc, block->getArgument(blockArgIndex).getType(),
+        loc, memref::getTensorTypeFromMemRefType(block->getArgument(blockArgIndex).getType()),
         block->getArgument(blockArgIndex), isRestrict, isWritable);
     ArrayRef<int64_t> origShape =
         cast<ShapedType>(argTypes[blockArgIndex]).getShape();
@@ -3792,10 +3796,10 @@ static func::FuncOp createCpuAttentionKernelWithMlir(ModuleOp module,
   }
 
   Value output = block->getArguments().back();
-  auto outputType = cast<MemRefType>(output.getType());
+  auto outputType = cast<mlir::bufferization::BufferLikeType>(output.getType());
   ImplicitLocOpBuilder implicitBuilder(loc, builder);
-  auto shapeValue =
-      tosa::getTosaConstShape(implicitBuilder, outputType.getShape());
+  auto shapeValue = tosa::getTosaConstShape(
+      implicitBuilder, cast<ShapedType>(outputType).getShape());
   auto flatResultTensor =
       builder.create<tosa::ReshapeOp>(loc, resultTensor, shapeValue);
 
@@ -3806,9 +3810,9 @@ static func::FuncOp createCpuAttentionKernelWithMlir(ModuleOp module,
 
   // return LSE (log-sum-exp)
   if (returnLSE) {
-    auto lseOutType = cast<MemRefType>(lseOut.getType());
-    auto lseShapeValue =
-        tosa::getTosaConstShape(implicitBuilder, lseOutType.getShape());
+    auto lseOutType = cast<bufferization::BufferLikeType>(lseOut.getType());
+    auto lseShapeValue = tosa::getTosaConstShape(
+        implicitBuilder, cast<ShapedType>(lseOutType).getShape());
     auto flatLseTensor =
         builder.create<tosa::ReshapeOp>(loc, lseTensor, lseShapeValue);