[mlir][spirv] Lower memref.reinterpret_cast

[Hardcode84]

Summary

For kernel SPIR-V, we are lowering memref to bare pointers, so reinterpret can be lowered to pointer, adjusted by offset value.

[kuhar]

The logic seems fine to me, just some stylistic issues.

@antiagainst, can you also take a look?

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

260

nit: add final

265–266

For consistency with the other code in this file, this function should be defined outside of the class definition.

272–273

nit: Use llvm::formatv for the diag message: rewriter.notifyMatchFailure(op, llvm::formatv(...)). You can find examples in other spirv patterns, so this is mostly for consistency with that code.

276

nit: auto doesn't improve readability here IMO, consider spelling out the type whenever it's not obvious based on the RHS

279–280

nit: Use converter->convertType<spirv::PointerType>(...)

282–283

also here

298

nit: Location loc

[antiagainst]

LGTM. Just a few more nits.

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

277

This assert is not necessary? The converter pretty much always exists there.

296

"failed to convert index type" to be clear.

311

Nit: I'd just do

auto acOp = rewriter.create<spirv::InBoundsPtrAccessChainOp>(loc, src, offsetValue, std::nullopt);
rewriter.replaceOp(op, acOp.getResult());

to make it indent nicer.

[Hardcode84]

rebase, review comments

[Hardcode84]

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

272–273

Actually, I don't see any formatv usage in this file but I see diag usage.

311

Changed to replaceOpWithNewOp

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

 public:
   using OpConversionPattern<memref::StoreOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+class ReinterpretCastPattern final
+    : public OpConversionPattern<memref::ReinterpretCastOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::ReinterpretCastOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // AllocaOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 AllocaOpPattern::matchAndRewrite(memref::AllocaOp allocaOp, OpAdaptor adaptor,
                                  ConversionPatternRewriter &rewriter) const {
   MemRefType allocType = allocaOp.getType();
   if (!isAllocationSupported(allocaOp, allocType))
     return rewriter.notifyMatchFailure(allocaOp, "unhandled allocation type");
 
   // Get the SPIR-V type for the allocation.
   Type spirvType = getTypeConverter()->convertType(allocType);
   rewriter.replaceOpWithNewOp<spirv::VariableOp>(allocaOp, spirvType,
                                                  spirv::StorageClass::Function,
                                                  /*initializer=*/nullptr);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // AllocOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 AllocOpPattern::matchAndRewrite(memref::AllocOp operation, OpAdaptor adaptor,
                                 ConversionPatternRewriter &rewriter) const {
   MemRefType allocType = operation.getType();
   if (!isAllocationSupported(operation, allocType))
     return rewriter.notifyMatchFailure(operation, "unhandled allocation type");
 
   // Get the SPIR-V type for the allocation.
   Type spirvType = getTypeConverter()->convertType(allocType);
 
   // Insert spirv.GlobalVariable for this allocation.
   Operation *parent =
       SymbolTable::getNearestSymbolTable(operation->getParentOp());
   if (!parent)
     return failure();
   Location loc = operation.getLoc();
   spirv::GlobalVariableOp varOp;
   {
     OpBuilder::InsertionGuard guard(rewriter);
     Block &entryBlock = *parent->getRegion(0).begin();
     rewriter.setInsertionPointToStart(&entryBlock);
     auto varOps = entryBlock.getOps<spirv::GlobalVariableOp>();
     std::string varName =
         std::string("__workgroup_mem__") +
         std::to_string(std::distance(varOps.begin(), varOps.end()));
StoreOpPattern::matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
   if (!storePtr)
     return failure();
 
   rewriter.replaceOpWithNewOp<spirv::StoreOp>(storeOp, storePtr,
                                               adaptor.getValue());
   return success();
 }
 
+LogicalResult ReinterpretCastPattern::matchAndRewrite(
+    memref::ReinterpretCastOp op, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  Value src = adaptor.getSource();
+  auto srcType = dyn_cast<spirv::PointerType>(src.getType());
+
+  if (!srcType)
+    return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
+      diag << "invalid src type " << src.getType();
+    });
+
+  TypeConverter *converter = getTypeConverter();
+
+  auto dstType = converter->convertType<spirv::PointerType>(op.getType());
+  if (dstType != srcType)
+    return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
+      diag << "invalid dst type " << op.getType();
+    });
+
+  OpFoldResult offset =
+      getMixedValues(adaptor.getStaticOffsets(), adaptor.getOffsets(), rewriter)
+          .front();
+  if (isConstantIntValue(offset, 0)) {
+    rewriter.replaceOp(op, src);
+    return success();
+  }
+
+  Type intType = converter->convertType(rewriter.getIndexType());
+  if (!intType)
+    return rewriter.notifyMatchFailure(op, "failed to convert index type");
+
+  Location loc = op.getLoc();
+  auto offsetValue = [&]() -> Value {
+    if (auto val = dyn_cast<Value>(offset))
+      return val;
+
+    int64_t attrVal = cast<IntegerAttr>(offset.get<Attribute>()).getInt();
+    Attribute attr = rewriter.getIntegerAttr(intType, attrVal);
+    return rewriter.create<spirv::ConstantOp>(loc, intType, attr);
+  }();
+
+  rewriter.replaceOpWithNewOp<spirv::InBoundsPtrAccessChainOp>(
+      op, src, offsetValue, std::nullopt);
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Pattern population
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 void populateMemRefToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                    RewritePatternSet &patterns) {
-  patterns.add<AllocaOpPattern, AllocOpPattern, AtomicRMWOpPattern,
-               DeallocOpPattern, IntLoadOpPattern, IntStoreOpPattern,
-               LoadOpPattern, MemorySpaceCastOpPattern, StoreOpPattern>(
+  patterns
+      .add<AllocaOpPattern, AllocOpPattern, AtomicRMWOpPattern,
+           DeallocOpPattern, IntLoadOpPattern, IntStoreOpPattern, LoadOpPattern,
+           MemorySpaceCastOpPattern, StoreOpPattern, ReinterpretCastPattern>(
           typeConverter, patterns.getContext());
 }
 } // namespace mlir

mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir

 func.func @load_i16(%arg0: memref<10xi16, #spirv.storage_class<StorageBuffer>>, %index : index) -> i16 {
   //     CHECK: %[[RES:.+]] = spirv.UConvert %{{.+}} : i32 to i16
   //     CHECK: return %[[RES]]
   %0 = memref.load %arg0[%index] : memref<10xi16, #spirv.storage_class<StorageBuffer>>
   return %0: i16
 }
 
 } // end module
+
+// -----
+
+// Check reinterpret_casts
+
+module attributes {
+  spirv.target_env = #spirv.target_env<
+    #spirv.vce<v1.0,
+      [Kernel, Addresses, GenericPointer], []>, #spirv.resource_limits<>>
+} {
+
+// CHECK-LABEL: func.func @reinterpret_cast
+//  CHECK-SAME:  (%[[MEM:.*]]: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>, %[[OFF:.*]]: index)
+func.func @reinterpret_cast(%arg: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>, %arg1: index) -> memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>> {
+//       CHECK:  %[[MEM1:.*]] = builtin.unrealized_conversion_cast %[[MEM]] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to !spirv.ptr<f32, CrossWorkgroup>
+//       CHECK:  %[[OFF1:.*]] = builtin.unrealized_conversion_cast %[[OFF]] : index to i32
+//       CHECK:  %[[RET:.*]] = spirv.InBoundsPtrAccessChain %[[MEM1]][%[[OFF1]]] : !spirv.ptr<f32, CrossWorkgroup>, i32
+//       CHECK:  %[[RET1:.*]] = builtin.unrealized_conversion_cast %[[RET]] : !spirv.ptr<f32, CrossWorkgroup> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+//       CHECK:  return %[[RET1]]
+  %ret = memref.reinterpret_cast %arg to offset: [%arg1], sizes: [10], strides: [1] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+  return %ret : memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+}
+
+// CHECK-LABEL: func.func @reinterpret_cast_0
+//  CHECK-SAME:  (%[[MEM:.*]]: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>)
+func.func @reinterpret_cast_0(%arg: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>) -> memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>> {
+//       CHECK:  %[[MEM1:.*]] = builtin.unrealized_conversion_cast %[[MEM]] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to !spirv.ptr<f32, CrossWorkgroup>
+//       CHECK:  %[[RET:.*]] = builtin.unrealized_conversion_cast %[[MEM1]] : !spirv.ptr<f32, CrossWorkgroup> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+//       CHECK:  return %[[RET]]
+  %ret = memref.reinterpret_cast %arg to offset: [0], sizes: [10], strides: [1] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+  return %ret : memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+}
+
+// CHECK-LABEL: func.func @reinterpret_cast_5
+//  CHECK-SAME:  (%[[MEM:.*]]: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>)
+func.func @reinterpret_cast_5(%arg: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>) -> memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>> {
+//       CHECK:  %[[MEM1:.*]] = builtin.unrealized_conversion_cast %[[MEM]] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to !spirv.ptr<f32, CrossWorkgroup>
+//       CHECK:  %[[OFF:.*]] = spirv.Constant 5 : i32
+//       CHECK:  %[[RET:.*]] = spirv.InBoundsPtrAccessChain %[[MEM1]][%[[OFF]]] : !spirv.ptr<f32, CrossWorkgroup>, i32
+//       CHECK:  %[[RET1:.*]] = builtin.unrealized_conversion_cast %[[RET]] : !spirv.ptr<f32, CrossWorkgroup> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+//       CHECK:  return %[[RET1]]
+  %ret = memref.reinterpret_cast %arg to offset: [5], sizes: [10], strides: [1] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+  return %ret : memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+}
+
+} // end module