[SYCL-MLIR] Add sycl.accessor.get_pointer (#8703)

This operation returns a pointer to the start of this accessor's memory.

---------

Signed-off-by: Tsang, Whitney <[email protected]>

Author: whitneywhtsang

mlir-sycl/include/mlir/Dialect/SYCL/IR/SYCLOps.td

@@ -668,6 +668,29 @@ def SYCLCallOp : SYCL_Op<"call", [CallOpInterface]> {
   }];
 }
 
+////////////////////////////////////////////////////////////////////////////////
+// accessor.get_pointer OPERATION
+////////////////////////////////////////////////////////////////////////////////
+
+def SYCLAccessorGetPointerOp
+    : SYCLMethodOpInterfaceImpl<"accessor.get_pointer", "AccessorType",
+                                ["get_pointer"]> {
+  let summary = "Represents the accessor get_pointer operation";
+  let description = [{
+    Returns a pointer to the start of this accessor's memory.
+  }];
+
+  let arguments = (ins Arg<AccessorMemRef, "The accessor", [MemRead]>:$Acc,
+                       TypeArrayAttr:$ArgumentTypes,
+                       FlatSymbolRefAttr:$FunctionName,
+                       OptionalAttr<FlatSymbolRefAttr>:$MangledFunctionName,
+                       FlatSymbolRefAttr:$TypeName);
+
+  let hasVerifier = 1;
+
+  let results = (outs AnyMemRef:$result);
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // accessor.size OPERATION
 ////////////////////////////////////////////////////////////////////////////////

mlir-sycl/lib/Conversion/SYCLToLLVM/SYCLToLLVM.cpp

@@ -133,6 +133,11 @@ struct AccessorGetPtr : public OffsetTag {
   static constexpr std::array<int32_t, 2> indices{1, 0};
 };
 
+/// Get the ID field from an accessor.
+struct AccessorGetID : public OffsetTag {
+  static constexpr std::array<int32_t, 2> indices{0, 0};
+};
+
 /// Get the MAccessRange field from an accessor.
 struct AccessorGetMAccessRange : public OffsetTag {
   static constexpr std::array<int32_t, 2> indices{0, 1};
@@ -1061,6 +1066,67 @@ class ConstructorPattern final
   }
 };
 
+//===----------------------------------------------------------------------===//
+// AccessorGetPointerPattern - Convert `sycl.accessor.get_pointer` to LLVM.
+//===----------------------------------------------------------------------===//
+
+class AccessorGetPointerPattern
+    : public ConvertOpToLLVMPattern<SYCLAccessorGetPointerOp>,
+      public GetMemberPattern<AccessorGetPtr>,
+      public GetMemberPattern<AccessorGetID, IDGetDim>,
+      public GetMemberPattern<AccessorGetMemRange, RangeGetDim> {
+public:
+  using ConvertOpToLLVMPattern<
+      SYCLAccessorGetPointerOp>::ConvertOpToLLVMPattern;
+
+private:
+  template <typename... Args> Value getID(Args &&...args) const {
+    return GetMemberPattern<AccessorGetID, IDGetDim>::loadValue(
+        std::forward<Args>(args)...);
+  }
+  template <typename... Args> Value getMemRange(Args &&...args) const {
+    return GetMemberPattern<AccessorGetMemRange, RangeGetDim>::loadValue(
+        std::forward<Args>(args)...);
+  }
+
+  Value getTotalOffset(OpBuilder &builder, Location loc, AccessorType accTy,
+                       OpAdaptor opAdaptor) const {
+    const auto acc = opAdaptor.getAcc();
+    const auto resTy = builder.getI64Type();
+    Value res = builder.create<arith::ConstantIntOp>(loc, 0, resTy);
+    for (unsigned i = 0; i < accTy.getDimension(); ++i) {
+      // Res = Res * Mem[I] + Id[I]
+      const auto memI = getMemRange(builder, loc, resTy, acc, i);
+      const auto idI = getID(builder, loc, resTy, acc, i);
+      res = builder.create<arith::AddIOp>(
+          loc, builder.create<arith::MulIOp>(loc, res, memI), idI);
+    }
+    return res;
+  }
+
+public:
+  LogicalResult
+  matchAndRewrite(SYCLAccessorGetPointerOp op, OpAdaptor opAdaptor,
+                  ConversionPatternRewriter &rewriter) const final {
+    const auto loc = op.getLoc();
+    const Value zero = rewriter.create<arith::ConstantIntOp>(loc, 0, 64);
+    Value index = rewriter.create<arith::SubIOp>(
+        loc, zero,
+        getTotalOffset(
+            rewriter, loc,
+            op.getAcc().getType().getElementType().cast<AccessorType>(),
+            opAdaptor));
+    const auto ptrTy = getTypeConverter()
+                           ->convertType(op.getType())
+                           .cast<LLVM::LLVMPointerType>();
+    Value ptr = GetMemberPattern<AccessorGetPtr>::loadValue(
+        rewriter, loc, ptrTy, opAdaptor.getAcc());
+    rewriter.replaceOpWithNewOp<LLVM::GEPOp>(op, ptrTy, ptr, index,
+                                             /*inbounds*/ true);
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // AccessorSizePattern - Convert `sycl.accessor.size` to LLVM.
 //===----------------------------------------------------------------------===//
@@ -2123,13 +2189,14 @@ void mlir::populateSYCLToLLVMConversionPatterns(
   patterns.add<CastPattern>(typeConverter);
   patterns.add<BarePtrCastPattern>(typeConverter, /*benefit*/ 2);
   patterns
-      .add<AccessorSizePattern, AddZeroArgPattern<SYCLIDGetOp>,
-           AddZeroArgPattern<SYCLItemGetIDOp>, AtomicSubscriptIDOffset,
-           BarePtrAddrSpaceCastPattern, GroupGetGroupIDPattern,
-           GroupGetGroupLinearRangePattern, GroupGetGroupRangeDimPattern,
-           GroupGetLocalIDPattern, GroupGetLocalLinearRangePattern,
-           GroupGetLocalRangeDimPattern, IDGetPattern, IDGetRefPattern,
-           ItemGetIDDimPattern, ItemGetRangeDimPattern, ItemGetRangePattern,
+      .add<AccessorGetPointerPattern, AccessorSizePattern,
+           AddZeroArgPattern<SYCLIDGetOp>, AddZeroArgPattern<SYCLItemGetIDOp>,
+           AtomicSubscriptIDOffset, BarePtrAddrSpaceCastPattern,
+           GroupGetGroupIDPattern, GroupGetGroupLinearRangePattern,
+           GroupGetGroupRangeDimPattern, GroupGetLocalIDPattern,
+           GroupGetLocalLinearRangePattern, GroupGetLocalRangeDimPattern,
+           IDGetPattern, IDGetRefPattern, ItemGetIDDimPattern,
+           ItemGetRangeDimPattern, ItemGetRangePattern,
            NDItemGetGlobalIDDimPattern, NDItemGetGlobalIDPattern,
            NDItemGetGroupPattern, NDItemGetGroupRangeDimPattern,
            NDItemGetLocalIDDimPattern, NDItemGetLocalLinearIDPattern,

mlir-sycl/lib/Dialect/IR/SYCLOps.cpp

@@ -89,6 +89,18 @@ bool SYCLAddrSpaceCastOp::areCastCompatible(TypeRange inputs,
           (outputMS == genericAddressSpace));
 }
 
+LogicalResult SYCLAccessorGetPointerOp::verify() {
+  const auto accTy = cast<AccessorType>(
+      cast<MemRefType>(getOperand().getType()).getElementType());
+  const Type resTy = getResult().getType();
+  const Type resElemTy = cast<MemRefType>(resTy).getElementType();
+  return (resElemTy != accTy.getType())
+             ? emitOpError(
+                   "Expecting a reference to this accessor's value type (")
+                   << accTy.getType() << "). Got " << resTy
+             : success();
+}
+
 LogicalResult SYCLAccessorSubscriptOp::verify() {
   // Available only when: (Dimensions > 0)
   // reference operator[](id<Dimensions> index) const;
@@ -98,11 +110,8 @@ LogicalResult SYCLAccessorSubscriptOp::verify() {
 
   // Available only when: (AccessMode != access_mode::atomic && Dimensions == 1)
   // reference operator[](size_t index) const;
-  const auto AccessorTy = getOperand(0)
-                              .getType()
-                              .cast<MemRefType>()
-                              .getElementType()
-                              .cast<AccessorType>();
+  const auto AccessorTy = cast<AccessorType>(
+      cast<MemRefType>(getOperand(0).getType()).getElementType());
 
   const unsigned Dimensions = AccessorTy.getDimension();
   if (Dimensions == 0)

mlir-sycl/test/Conversion/SYCLToLLVM/sycl-methods-to-llvm.mlir

@@ -155,6 +155,38 @@ func.func @test(%id: memref<?x!sycl_id_3_>, %idx: i32) -> memref<?xi64> {
 
 // -----
 
+//===-------------------------------------------------------------------------------------------------===//
+// sycl.accessor.get_pointer
+//===-------------------------------------------------------------------------------------------------===//
+
+!sycl_id_1_ = !sycl.id<[1], (!sycl.array<[1], (memref<1xi64>)>)>
+!sycl_range_1_ = !sycl.range<[1], (!sycl.array<[1], (memref<1xi64>)>)>
+!sycl_accessor_impl_device_1_ = !sycl.accessor_impl_device<[1], (!sycl_id_1_, !sycl_range_1_, !sycl_range_1_)>
+!sycl_accessor_1_i32_rw_gb = !sycl.accessor<[1, i32, read_write, global_buffer], (!sycl_accessor_impl_device_1_, !llvm.struct<(ptr<i32, 1>)>)>
+
+// CHECK-LABEL:   llvm.func @test(
+// CHECK-SAME:                    %[[VAL_0:.*]]: !llvm.ptr<[[ACCESSOR1:.*]]>) -> !llvm.ptr<i32, 1> {
+// CHECK-NEXT:      %0 = llvm.mlir.constant(0 : i64) : i64
+// CHECK-NEXT:      %1 = llvm.mlir.constant(0 : i64) : i64
+// CHECK-NEXT:      %2 = llvm.getelementptr inbounds %arg0[0, 0, 2, 0, 0, 0] : (!llvm.ptr<[[ACCESSOR1]]>) -> !llvm.ptr<i64>
+// CHECK-NEXT:      %3 = llvm.load %2 : !llvm.ptr<i64>
+// CHECK-NEXT:      %4 = llvm.getelementptr inbounds %arg0[0, 0, 0, 0, 0, 0] : (!llvm.ptr<[[ACCESSOR1]]>) -> !llvm.ptr<i64>
+// CHECK-NEXT:      %5 = llvm.load %4 : !llvm.ptr<i64>
+// CHECK-NEXT:      %6 = llvm.mul %1, %3  : i64
+// CHECK-NEXT:      %7 = llvm.add %6, %5  : i64
+// CHECK-NEXT:      %8 = llvm.sub %0, %7  : i64
+// CHECK-NEXT:      %9 = llvm.getelementptr inbounds %arg0[0, 1, 0] : (!llvm.ptr<[[ACCESSOR1]]>) -> !llvm.ptr<ptr<i32, 1>>
+// CHECK-NEXT:      %10 = llvm.load %9 : !llvm.ptr<ptr<i32, 1>>
+// CHECK-NEXT:      %11 = llvm.getelementptr inbounds %10[%8] : (!llvm.ptr<i32, 1>, i64) -> !llvm.ptr<i32, 1>
+// CHECK-NEXT:      llvm.return %11 : !llvm.ptr<i32, 1>
+// CHECK-NEXT:    }
+func.func @test(%acc: memref<?x!sycl_accessor_1_i32_rw_gb>) -> memref<?xi32, 1> {
+  %0 = sycl.accessor.get_pointer(%acc) { ArgumentTypes = [memref<?x!sycl_accessor_1_i32_rw_gb>], FunctionName = @"get_pointer", MangledFunctionName = @"get_pointer", TypeName = @"accessor" }  : (memref<?x!sycl_accessor_1_i32_rw_gb>) -> memref<?xi32, 1>
+  return %0 : memref<?xi32, 1>
+}
+
+// -----
+
 //===-------------------------------------------------------------------------------------------------===//
 // sycl.accessor.size
 //===-------------------------------------------------------------------------------------------------===//

mlir-sycl/test/Dialect/IR/SYCL/invalid.mlir

@@ -260,6 +260,19 @@ func.func @test_work_group_id_dim() -> index {
 
 // -----
 
+!sycl_id_1_ = !sycl.id<[1], (!sycl.array<[1], (memref<1xi64>)>)>
+!sycl_range_1_ = !sycl.range<[1], (!sycl.array<[1], (memref<1xi64>)>)>
+!sycl_accessor_impl_device_1_ = !sycl.accessor_impl_device<[1], (!sycl_id_1_, !sycl_range_1_, !sycl_range_1_)>
+!sycl_accessor_1_i32_rw_gb = !sycl.accessor<[1, i32, read_write, global_buffer], (!sycl_accessor_impl_device_1_, !llvm.struct<(ptr<i32, 1>)>)>
+
+func.func @test_accessor_get_pointer(%acc: memref<?x!sycl_accessor_1_i32_rw_gb>) -> memref<?xi64, 1> {
+  // expected-error @+1 {{'sycl.accessor.get_pointer' op Expecting a reference to this accessor's value type}}
+  %0 = sycl.accessor.get_pointer(%acc) { ArgumentTypes = [memref<?x!sycl_accessor_1_i32_rw_gb>], FunctionName = @"get_pointer", MangledFunctionName = @"get_pointer", TypeName = @"accessor" }  : (memref<?x!sycl_accessor_1_i32_rw_gb>) -> memref<?xi64, 1>
+  return %0 : memref<?xi64, 1>
+}
+
+// -----
+
 !sycl_id_1_ = !sycl.id<[1], (!sycl.array<[1], (memref<1xi64, 4>)>)>
 !sycl_range_1_ = !sycl.range<[1], (!sycl.array<[1], (memref<1xi64, 4>)>)>
 !sycl_accessor_impl_device_1_ = !sycl.accessor_impl_device<[1], (!sycl_id_1_, !sycl_range_1_, !sycl_range_1_)>

mlir-sycl/test/Dialect/IR/SYCL/ops.mlir

@@ -260,6 +260,12 @@ func.func @test_work_group_id_const() -> index {
   return %0 : index
 }
 
+// CHECL-LABEL: test_accessor_get_pointer
+func.func @test_accessor_get_pointer(%acc: memref<?x!sycl_accessor_1_i32_w_gb>) -> memref<?xi32, 1> {
+  %0 = sycl.accessor.get_pointer(%acc) { ArgumentTypes = [memref<?x!sycl_accessor_1_i32_w_gb>], FunctionName = @"get_pointer", MangledFunctionName = @"get_pointer", TypeName = @"accessor" }  : (memref<?x!sycl_accessor_1_i32_w_gb>) -> memref<?xi32, 1>
+  return %0 : memref<?xi32, 1>
+}
+
 // CHECK-LABEL: test_accessor_subscript_atomic
 func.func @test_accessor_subscript_atomic(
   %acc: memref<?x!sycl_accessor_1_i32_ato_gb>, 

polygeist/tools/cgeist/Test/Verification/sycl/functions.cpp

@@ -40,6 +40,23 @@
 
 template <typename T> SYCL_EXTERNAL void keep(T);
 
+// COM: Commenting out the checks below, this is the code that should be
+// generated. Currently the DPC++ SYCL RT implementation of
+// accessor::get_pointer is non-conforming. Once that problem is fixed we
+// enable the checks below (with the MangledFunctionName fixed).
+
+// COM-MLIR-LABEL: func.func @_Z20accessor_get_pointerN4sycl3_V18accessorIiLi2ELNS0_6access4modeE1026ELNS2_6targetE2014ELNS2_11placeholderE0ENS0_3ext6oneapi22accessor_property_listIJEEEEE(
+// COM-MLIR:           %{{.*}}: memref<?x!sycl_accessor_2_i32_rw_gb> {llvm.align = 8 : i64, llvm.byval = !sycl_accessor_2_i32_rw_gb, llvm.noundef})
+// COM-MLIR: %{{.*}} = sycl.accessor.get_pointer(%{{.*}}) {ArgumentTypes = [memref<?x!sycl_accessor_2_i32_rw_gb, 4>], FunctionName = @get_pointer, MangledFunctionName = @_ZNK4sycl3_V18accessorIiLi2ELNS0_6access4modeE1026ELNS2_6targetE2014ELNS2_11placeholderE0ENS0_3ext6oneapi22accessor_property_listIJEEEE11get_pointerILS4_2014EvEENS0_9IiLNS2_13address_spaceE1ELNS2_9decoratedE2EEEv, TypeName = @accessor} : (memref<?x!sycl_accessor_2_i32_rw_gb>) -> memref<?xi32, 1>
+
+// COM-LLVM-LABEL: define spir_func void @_Z20accessor_get_pointerN4sycl3_V18accessorIiLi2ELNS0_6access4modeE1026ELNS2_6targetE2014ELNS2_11placeholderE0ENS0_3ext6oneapi22accessor_property_listIJEEEEE(
+// COM-LLVM:           %"class.sycl::_V1::accessor.2"* noundef byval(%"class.sycl::_V1::accessor.2") align 8 %0) #[[FUNCATTRS:[0-9]+]] {
+// COM-LLVM:  %{{.*}} = call spir_func i32 addrspace(1)* @_ZNK4sycl3_V18accessorIiLi2ELNS0_6access4modeE1026ELNS2_6targetE2014ELNS2_11placeholderE0ENS0_3ext6oneapi22accessor_property_listIJEEEE11get_pointerILS4_2014EvEENS0_9IiLNS2_13address_spaceE1ELNS2_9decoratedE2EEEv(%"class.sycl::_V1::accessor.2" addrspace(4)* %{{.*}})
+
+SYCL_EXTERNAL void accessor_get_pointer(sycl::accessor<sycl::cl_int, 2> acc) {
+  keep(acc.get_pointer());
+}
+
 // CHECK-MLIR-LABEL: func.func @_Z13accessor_sizeN4sycl3_V18accessorIiLi2ELNS0_6access4modeE1026ELNS2_6targetE2014ELNS2_11placeholderE0ENS0_3ext6oneapi22accessor_property_listIJEEEEE(
 // CHECK-MLIR:           %{{.*}}: memref<?x!sycl_accessor_2_i32_rw_gb> {llvm.align = 8 : i64, llvm.byval = !sycl_accessor_2_i32_rw_gb, llvm.noundef})
 // CHECK-MLIR: %{{.*}} = sycl.accessor.size(%arg0) {ArgumentTypes = [memref<?x!sycl_accessor_2_i32_rw_gb, 4>], FunctionName = @size, MangledFunctionName = @_ZNK4sycl3_V18accessorIiLi2ELNS0_6access4modeE1026ELNS2_6targetE2014ELNS2_11placeholderE0ENS0_3ext6oneapi22accessor_property_listIJEEEE4sizeEv, TypeName = @accessor} : (memref<?x!sycl_accessor_2_i32_rw_gb>) -> i64