Conversion of 'sycl.constructor(%0, %1) {type = @id}' to a copy ctor. (#50)

This PR adds support for converting a `sycl.constructor(%0, %1) {type = @id}` operation (representing a copy construction) to a call to the `sycl::id<n>` copy constructor (`sycl::id<n>::id(const sycl::id<n> const&)`).

Signed-off-by: Tiotto, Ettore <[email protected]>

Author: etiotto

mlir-sycl/include/mlir/Conversion/SYCLToLLVM/SYCLFuncRegistry.h

@@ -35,6 +35,8 @@ class SYCLFuncDescriptor {
   /// Enumerates SYCL functions.
   // clang-format off
   enum class FuncId {
+    Unknown, 
+
     // Member functions for the sycl:id<n> class.
     Id1CtorDefault, // sycl::id<1>::id()
     Id2CtorDefault, // sycl::id<2>::id()
@@ -48,11 +50,20 @@ class SYCLFuncDescriptor {
     Id1CtorItem,    // sycl::id<1>::id<1>(std::enable_if<(1)==(1), unsigned long>::type, unsigned long, unsigned long)
     Id2CtorItem,    // sycl::id<2>::id<2>(std::enable_if<(2)==(2), unsigned long>::type, unsigned long, unsigned long)
     Id3CtorItem,    // sycl::id<3>::id<3>(std::enable_if<(3)==(3), unsigned long>::type, unsigned long, unsigned long)
-
+    Id1CopyCtor,    // sycl::id<1>::id(sycl::id<1> const&)
+    Id2CopyCtor,    // sycl::id<2>::id(sycl::id<2> const&)
+    Id3CopyCtor,    // sycl::id<3>::id(sycl::id<3> const&)
+    
     // Member functions for ..TODO..
   };
   // clang-format on
 
+  /// Enumerates the kind of FuncId.
+  enum class FuncIdKind {
+    Unknown,
+    IdCtor, // any sycl::id<n> constructors
+  };
+
   // Call the SYCL constructor identified by \p id with the given \p args.
   static Value call(FuncId id, ValueRange args,
                     const SYCLFuncRegistry &registry, OpBuilder &b,
@@ -65,9 +76,12 @@ class SYCLFuncDescriptor {
       : id(id), name(name), outputTy(outputTy),
         argTys(argTys.begin(), argTys.end()) {}
 
-  // Inject the declaration for this function into the module.
+  /// Inject the declaration for this function into the module.
   void declareFunction(ModuleOp &module, OpBuilder &b);
 
+  /// Returns true if the given \p funcId is for a sycl::id<n> constructor.
+  static bool isIdCtor(FuncId funcId);
+
 private:
   FuncId id;                   // unique identifier for a SYCL function
   StringRef name;              // SYCL function name
@@ -86,12 +100,18 @@ class SYCLFuncRegistry {
   static const SYCLFuncRegistry create(ModuleOp &module, OpBuilder &builder);
 
   const SYCLFuncDescriptor &getFuncDesc(SYCLFuncDescriptor::FuncId id) const {
-    assert(
-        (registry.find(id) != registry.end()) &&
-        "function identified by 'id' not found in the SYCL function registry");
+    assert((registry.find(id) != registry.end()) &&
+           "function identified by 'id' not found in the SYCL function "
+           "registry");
     return registry.at(id);
   }
 
+  // Returns the SYCLFuncDescriptor::FuncId corresponding to the function
+  // descriptor that matches the given signature and funcIdKind.
+  SYCLFuncDescriptor::FuncId
+  getFuncId(SYCLFuncDescriptor::FuncIdKind funcIdKind, Type retType,
+            TypeRange argTypes) const;
+
 private:
   SYCLFuncRegistry(ModuleOp &module, OpBuilder &builder);
 

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

@@ -32,10 +32,34 @@ void SYCLFuncDescriptor::declareFunction(ModuleOp &module, OpBuilder &b) {
   funcRef = builder.getOrInsertFuncDecl(name, outputTy, argTys, module);
 }
 
-Value SYCLFuncDescriptor::call(FuncId id, ValueRange args,
+bool SYCLFuncDescriptor::isIdCtor(FuncId funcId) {
+  switch (funcId) {
+  case FuncId::Id1CtorDefault:
+  case FuncId::Id2CtorDefault:
+  case FuncId::Id3CtorDefault:
+  case FuncId::Id1CtorSizeT:
+  case FuncId::Id2CtorSizeT:
+  case FuncId::Id3CtorSizeT:
+  case FuncId::Id1CtorRange:
+  case FuncId::Id2CtorRange:
+  case FuncId::Id3CtorRange:
+  case FuncId::Id1CtorItem:
+  case FuncId::Id2CtorItem:
+  case FuncId::Id3CtorItem:
+  case FuncId::Id1CopyCtor:
+  case FuncId::Id2CopyCtor:
+  case FuncId::Id3CopyCtor:
+    return true;
+  default:;
+  }
+
+  return false;
+}
+
+Value SYCLFuncDescriptor::call(FuncId funcId, ValueRange args,
                                const SYCLFuncRegistry &registry, OpBuilder &b,
                                Location loc) {
-  const SYCLFuncDescriptor &funcDesc = registry.getFuncDesc(id);
+  const SYCLFuncDescriptor &funcDesc = registry.getFuncDesc(funcId);
   LLVM_DEBUG(
       llvm::dbgs() << "Creating SYCLFuncDescriptor::call to funcDesc.funcRef: "
                    << funcDesc.funcRef << "\n");
@@ -59,14 +83,56 @@ Value SYCLFuncDescriptor::call(FuncId id, ValueRange args,
 
 SYCLFuncRegistry *SYCLFuncRegistry::instance = nullptr;
 
-const SYCLFuncRegistry SYCLFuncRegistry::create(
-    ModuleOp &module, OpBuilder &builder) {
+const SYCLFuncRegistry SYCLFuncRegistry::create(ModuleOp &module,
+                                                OpBuilder &builder) {
   if (!instance)
     instance = new SYCLFuncRegistry(module, builder);
 
   return *instance;
 }
 
+SYCLFuncDescriptor::FuncId
+SYCLFuncRegistry::getFuncId(SYCLFuncDescriptor::FuncIdKind funcIdKind,
+                            Type retType, TypeRange argTypes) const {
+  assert(funcIdKind != SYCLFuncDescriptor::FuncIdKind::Unknown &&
+         "Invalid funcIdKind");
+
+  // Determines whether the given funcId has kind that matches the given
+  // funcIdKind.
+  auto kindMatches = [](SYCLFuncDescriptor::FuncId funcId,
+                        SYCLFuncDescriptor::FuncIdKind funcIdKind) {
+    bool foundMatch = false;
+    switch (funcIdKind) {
+    case SYCLFuncDescriptor::FuncIdKind::IdCtor:
+      foundMatch = SYCLFuncDescriptor::isIdCtor(funcId);
+      break;
+    default:
+      foundMatch = false;
+    }
+    return foundMatch;
+  };
+
+  for (const auto &entry : registry) {
+    // Skip through entries that do not match the requested funcIdKind.
+    if (!kindMatches(entry.second.id, funcIdKind))
+      continue;
+
+    // Ensure that the entry has return and arguments type that match the one
+    // provided.
+    if (retType != entry.second.outputTy ||
+        argTypes.size() != entry.second.argTys.size())
+      continue;
+    if (!std::equal(argTypes.begin(), argTypes.end(),
+                    entry.second.argTys.begin()))
+      continue;
+
+    return entry.second.id;
+  }
+
+  llvm_unreachable("Unimplemented descriptor");
+  return SYCLFuncDescriptor::FuncId::Unknown;
+}
+
 SYCLFuncRegistry::SYCLFuncRegistry(ModuleOp &module, OpBuilder &builder)
     : registry() {
   MLIRContext *context = module.getContext();
@@ -144,6 +210,22 @@ SYCLFuncRegistry::SYCLFuncRegistry(ModuleOp &module, OpBuilder &builder)
           SYCLFuncDescriptor::FuncId::Id3CtorItem,
           "_ZN2cl4sycl2idILi3EEC2ILi3EEENSt9enable_ifIXeqT_Li3EEmE4typeEmm",
           voidTy, {id3PtrTy, i64Ty, i64Ty, i64Ty}),
+
+      // cl::sycl::id<1>::id(cl::sycl::id<1> const&)
+      SYCLFuncDescriptor(
+          SYCLFuncDescriptor::FuncId::Id1CopyCtor,
+          "_ZN2cl4sycl2idILi1EEC1ERKS2_",
+          voidTy, {id1PtrTy, id1PtrTy}),
+      // cl::sycl::id<2>::id(cl::sycl::id<2> const&)
+      SYCLFuncDescriptor(
+          SYCLFuncDescriptor::FuncId::Id2CopyCtor,
+          "_ZN2cl4sycl2idILi2EEC1ERKS2_",
+          voidTy, {id2PtrTy, id2PtrTy}),
+      // cl::sycl::id<3>::id(cl::sycl::id<3> const&)
+      SYCLFuncDescriptor(
+          SYCLFuncDescriptor::FuncId::Id3CopyCtor,
+          "_ZN2cl4sycl2idILi3EEC1ERKS2_",
+          voidTy, {id3PtrTy, id3PtrTy}),
   };
   // clang-format on
 

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

@@ -185,32 +185,14 @@ class ConstructorPattern final
                llvm::dbgs() << "\n");
 
     ModuleOp module = op.getOperation()->getParentOfType<ModuleOp>();
-    const auto &registry = SYCLFuncRegistry::create(module, rewriter);
-    ValueRange args(op.Args());
-    assert(args.size() > 0 && "Expecting at least one argument (the this ptr)");
-
-    // Multikey map used to lookup the specific sycl::id's ctor to call.
-    // Given 'sycl::id<dim>(args)' the key is the pair {size(args), dim}.
-    const std::map<std::pair<int, int>, SYCLFuncDescriptor::FuncId>
-        lookupCtorId = {
-            {{1, 1}, SYCLFuncDescriptor::FuncId::Id1CtorDefault},
-            {{1, 2}, SYCLFuncDescriptor::FuncId::Id2CtorDefault},
-            {{1, 3}, SYCLFuncDescriptor::FuncId::Id3CtorDefault},
-            {{2, 1}, SYCLFuncDescriptor::FuncId::Id1CtorSizeT},
-            {{2, 2}, SYCLFuncDescriptor::FuncId::Id2CtorSizeT},
-            {{2, 3}, SYCLFuncDescriptor::FuncId::Id3CtorSizeT},
-            {{3, 1}, SYCLFuncDescriptor::FuncId::Id1CtorRange},
-            {{3, 2}, SYCLFuncDescriptor::FuncId::Id2CtorRange},
-            {{3, 3}, SYCLFuncDescriptor::FuncId::Id3CtorRange},
-            {{4, 1}, SYCLFuncDescriptor::FuncId::Id1CtorItem},
-            {{4, 2}, SYCLFuncDescriptor::FuncId::Id2CtorItem},
-            {{4, 3}, SYCLFuncDescriptor::FuncId::Id3CtorItem},
-        };
+    MLIRContext *context = module.getContext();
 
     // Lookup the ctor function to use.
-    auto arg0ElemTy = getElementType<mlir::sycl::IDType>(args[0].getType());   
-    auto key = std::make_pair(args.size(), arg0ElemTy.getDimension());
-    SYCLFuncDescriptor::FuncId funcId = lookupCtorId.at(key);
+    const auto &registry = SYCLFuncRegistry::create(module, rewriter);    
+    auto voidTy = LLVM::LLVMVoidType::get(context);
+    SYCLFuncDescriptor::FuncId funcId =
+        registry.getFuncId(SYCLFuncDescriptor::FuncIdKind::IdCtor, voidTy,
+                           opAdaptor.Args().getTypes());
 
     // Generate an LLVM call to the appropriate ctor.
     SYCLFuncDescriptor::call(funcId, opAdaptor.getOperands(), registry,

mlir-sycl/test/Conversion/SYCLToLLVM/func-ops-to-llvm.mlir

@@ -124,3 +124,34 @@ func.func @id3CtorItem(%arg0: memref<?x!sycl.id<3>>, %arg1: i64, %arg2: i64, %ar
 }
 
 // -----
+
+//===-------------------------------------------------------------------------------------------------===//
+// Constructors sycl::id<n>::id(sycl::id<n> const&)
+//===-------------------------------------------------------------------------------------------------===//
+
+// CHECK: llvm.func @_ZN2cl4sycl2idILi1EEC1ERKS2_([[THIS_PTR_TYPE:!llvm.struct<\(ptr<struct<"class.cl::sycl::id.1",.*]], [[THIS_PTR_TYPE]])
+func.func @id1CopyCtor(%arg0: memref<?x!sycl.id<1>>, %arg1: memref<?x!sycl.id<1>>) {
+  // CHECK: llvm.call @_ZN2cl4sycl2idILi1EEC1ERKS2_({{.*}}, {{.*}}) : ([[THIS_PTR_TYPE]], [[THIS_PTR_TYPE]]) -> ()  
+ "sycl.constructor"(%arg0, %arg1) {Type = @id} : (memref<?x!sycl.id<1>>, memref<?x!sycl.id<1>>) -> ()
+  return
+}
+
+// -----
+
+// CHECK: llvm.func @_ZN2cl4sycl2idILi2EEC1ERKS2_([[THIS_PTR_TYPE:!llvm.struct<\(ptr<struct<"class.cl::sycl::id.2",.*]], [[THIS_PTR_TYPE]])
+func.func @id2CopyCtor(%arg0: memref<?x!sycl.id<2>>, %arg1: memref<?x!sycl.id<2>>) {
+  // CHECK: llvm.call @_ZN2cl4sycl2idILi2EEC1ERKS2_({{.*}}, {{.*}}) : ([[THIS_PTR_TYPE]], [[THIS_PTR_TYPE]]) -> ()
+ "sycl.constructor"(%arg0, %arg1) {Type = @id} : (memref<?x!sycl.id<2>>, memref<?x!sycl.id<2>>) -> ()  
+  return
+}
+
+// -----
+
+// CHECK: llvm.func @_ZN2cl4sycl2idILi3EEC1ERKS2_([[THIS_PTR_TYPE:!llvm.struct<\(ptr<struct<"class.cl::sycl::id.3",.*]], [[THIS_PTR_TYPE]])
+func.func @id3CopyCtor(%arg0: memref<?x!sycl.id<3>>, %arg1: memref<?x!sycl.id<3>>) {
+  // CHECK: llvm.call @_ZN2cl4sycl2idILi3EEC1ERKS2_({{.*}}, {{.*}}) : ([[THIS_PTR_TYPE]], [[THIS_PTR_TYPE]]) -> ()
+ "sycl.constructor"(%arg0, %arg1) {Type = @id} : (memref<?x!sycl.id<3>>, memref<?x!sycl.id<3>>) -> ()
+  return
+}
+
+// -----