Fix rocmlir-gen device selection (#1964)

Author: djramic

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

@@ -1362,16 +1362,37 @@ static func::FuncOp createGPUWrapper(ModuleOp module,
       func::FuncOp::create(loc, StringRef(funcNameGpu), gpuWrapperFuncType);
   module.push_back(gpuWrapperFunc);
 
+  // Emit device selection
+  if (deviceNum.getNumOccurrences() > 0) {
+    const int32_t priority = 122;
+    const StringRef constructorName = "setDeviceCtor";
+    auto func = module.lookupSymbol<mlir::LLVM::LLVMFuncOp>(constructorName);
+    if (!func) {
+      func = b.create<mlir::LLVM::LLVMFuncOp>(
+          module.getLoc(), constructorName,
+          mlir::LLVM::LLVMFunctionType::get(LLVM::LLVMVoidType::get(context),
+                                            {}));
+      module.push_back(func);
+
+      Block *block = func.addEntryBlock(b);
+      b.setInsertionPoint(block, block->begin());
+      b.create<gpu::SetDefaultDeviceOp>(
+          loc, b.create<arith::ConstantIntOp>(loc, b.getIntegerType(32),
+                                              deviceNum.getValue()));
+      b.create<mlir::LLVM::ReturnOp>(loc, ValueRange{});
+
+      b.setInsertionPointToEnd(module.getBody());
+      b.create<mlir::LLVM::GlobalCtorsOp>(
+          loc, b.getArrayAttr(mlir::SymbolRefAttr::get(func)),
+          b.getI32ArrayAttr({priority}),
+          b.getArrayAttr(mlir::LLVM::ZeroAttr::get(context)));
+    }
+  }
+
   // Emit gpu convolution logic.
   Block *block = gpuWrapperFunc.addEntryBlock();
   b.setInsertionPoint(block, block->begin());
 
-  // Emit device selection
-  if (deviceNum.getNumOccurrences() > 0)
-    b.create<gpu::SetDefaultDeviceOp>(
-        loc, b.create<arith::ConstantIntOp>(loc, b.getIntegerType(32),
-                                            deviceNum.getValue()));
-
   SmallVector<Value, 4> cpuMem;
   SmallVector<Value, 4> gpuMem;
   for (auto pair : llvm::enumerate(kernels[0].params)) {
@@ -5006,7 +5027,8 @@ int main(int argc, char **argv) {
                       math::MathDialect, arith::ArithDialect,
                       vector::VectorDialect, gpu::GPUDialect,
                       linalg::LinalgDialect, mhal::MHALDialect,
-                      bufferization::BufferizationDialect, tosa::TosaDialect>();
+                      bufferization::BufferizationDialect, tosa::TosaDialect,
+                      mlir::LLVM::LLVMDialect>();
 
   // Parse pass names in main to ensure static initialization completed.
   llvm::cl::ParseCommandLineOptions(argc, argv,