Fix review comments.

Signed-off-by: Ilya Enkovich <[email protected]>

Author: ienkovich

mlir/lib/Dialect/AMX/Transforms/LegalizeForLLVMExport.cpp

@@ -37,44 +37,38 @@ std::pair<Value, Value> getTileSizes(ConversionPatternRewriter &rewriter,
       rewriter.create<LLVM::ConstantOp>(loc, llvmInt16Type, nattr));
 }
 
-/// Verifies if the stride matches proper tile access.
-LogicalResult verifyStride(MemRefType mType) {
-  if (mType.getRank() < 2)
-    return failure();
-  int64_t last = mType.getRank() - 1;
-  int64_t offset;
-  SmallVector<int64_t, 4> strides;
-  if (failed(getStridesAndOffset(mType, strides, offset)) || strides[last] != 1)
-    return failure();
-  return success();
-}
-
 /// Maps the 2-dim memref shape to the 64-bit stride. Note that the buffer
 /// shape may "envelop" the actual tile shape, and may be dynamically sized.
-Value getStride(ConversionPatternRewriter &rewriter,
-                const LLVMTypeConverter &typeConverter, MemRefType mType,
-                Value base, Location loc) {
-  assert(mType.getRank() >= 2);
+/// Returns failure if proper stride couldn't be found.
+FailureOr<Value> getStride(ConversionPatternRewriter &rewriter,
+                           const LLVMTypeConverter &typeConverter,
+                           MemRefType mType, Value base, Location loc) {
+  if (mType.getRank() < 2)
+    return failure();
   int64_t preLast = mType.getRank() - 2;
   Type llvmInt64Type = IntegerType::get(&typeConverter.getContext(), 64);
   unsigned width = mType.getElementType().getIntOrFloatBitWidth();
   assert(llvm::isPowerOf2_64(width) && width >= 8);
   unsigned bytes = width >> 3;
   int64_t offset;
   SmallVector<int64_t, 4> strides;
-  getStridesAndOffset(mType, strides, offset);
+  if (failed(getStridesAndOffset(mType, strides, offset)) ||
+      strides.back() != 1)
+    return failure();
   if (strides[preLast] == ShapedType::kDynamic) {
     // Dynamic stride needs code to compute the stride at runtime.
     MemRefDescriptor memrefDescriptor(base);
     auto attr = rewriter.getI64IntegerAttr(bytes);
     Value scale = rewriter.create<LLVM::ConstantOp>(loc, llvmInt64Type, attr);
-    return rewriter.create<LLVM::MulOp>(
-        loc, llvmInt64Type, scale,
-        memrefDescriptor.stride(rewriter, loc, preLast));
+    return rewriter
+        .create<LLVM::MulOp>(loc, llvmInt64Type, scale,
+                             memrefDescriptor.stride(rewriter, loc, preLast))
+        .getResult();
   }
   // Use direct constant for static stride.
   auto attr = rewriter.getI64IntegerAttr(strides[preLast] * bytes);
-  return rewriter.create<LLVM::ConstantOp>(loc, llvmInt64Type, attr);
+  return rewriter.create<LLVM::ConstantOp>(loc, llvmInt64Type, attr)
+      .getResult();
 }
 
 struct TileZeroConversion : public ConvertOpToLLVMPattern<TileZeroOp> {
@@ -106,16 +100,16 @@ struct TileLoadConversion : public ConvertOpToLLVMPattern<TileLoadOp> {
     std::pair<Value, Value> tsz =
         getTileSizes(rewriter, *getTypeConverter(), vType, op.getLoc());
     // Determine stride.
-    if (failed(verifyStride(mType)))
+    auto stride = getStride(rewriter, *getTypeConverter(), mType,
+                            adaptor.getBase(), op.getLoc());
+    if (failed(stride))
       return failure();
-    Value stride = getStride(rewriter, *getTypeConverter(), mType,
-                             adaptor.getBase(), op.getLoc());
     // Replace operation with intrinsic.
     Value ptr = getStridedElementPtr(op.getLoc(), mType, adaptor.getBase(),
                                      adaptor.getIndices(), rewriter);
     Type resType = typeConverter->convertType(vType);
     rewriter.replaceOpWithNewOp<amx::x86_amx_tileloadd64>(
-        op, resType, tsz.first, tsz.second, ptr, stride);
+        op, resType, tsz.first, tsz.second, ptr, stride.value());
     return success();
   }
 };
@@ -132,15 +126,15 @@ struct TileStoreConversion : public ConvertOpToLLVMPattern<TileStoreOp> {
     std::pair<Value, Value> tsz =
         getTileSizes(rewriter, *getTypeConverter(), vType, op.getLoc());
     // Determine stride.
-    if (failed(verifyStride(mType)))
+    auto stride = getStride(rewriter, *getTypeConverter(), mType,
+                            adaptor.getBase(), op.getLoc());
+    if (failed(stride))
       return failure();
-    Value stride = getStride(rewriter, *getTypeConverter(), mType,
-                             adaptor.getBase(), op.getLoc());
     // Replace operation with intrinsic.
     Value ptr = getStridedElementPtr(op.getLoc(), mType, adaptor.getBase(),
                                      adaptor.getIndices(), rewriter);
     rewriter.replaceOpWithNewOp<amx::x86_amx_tilestored64>(
-        op, tsz.first, tsz.second, ptr, stride, adaptor.getVal());
+        op, tsz.first, tsz.second, ptr, stride.value(), adaptor.getVal());
     return success();
   }
 };