Use correct graphics allocation in setArgBuffer()

Use the rootDeviceIndex to select the graphics allocation associated
with the argument being set.

Change-Id: I2298f46c0ce5d96841d17381afb7b0013a3f804e
Signed-off: Jaime Arteaga <jaime.a.arteaga.molina@intel.com>

level_zero/core/source/kernel/kernel_imp.cpp

@@ -496,7 +496,8 @@ ze_result_t KernelImp::setArgBuffer(uint32_t argIndex, size_t argSize, const voi
     if (nullptr == svmAllocsManager->getSVMAlloc(requestedAddress)) {
         return ZE_RESULT_ERROR_INVALID_ARGUMENT;
     }
-    NEO::GraphicsAllocation *alloc = svmAllocsManager->getSVMAlloc(requestedAddress)->gpuAllocations.getDefaultGraphicsAllocation();
+    uint32_t rootDeviceIndex = module->getDevice()->getRootDeviceIndex();
+    NEO::GraphicsAllocation *alloc = svmAllocsManager->getSVMAlloc(requestedAddress)->gpuAllocations.getGraphicsAllocation(rootDeviceIndex);
     auto gpuAddress = reinterpret_cast<uintptr_t>(requestedAddress);
     return setArgBufferWithAlloc(argIndex, gpuAddress, alloc);
 }

level_zero/core/test/unit_tests/fixtures/device_fixture.h

@@ -7,6 +7,8 @@
 
 #pragma once
 
+#include "shared/source/os_interface/device_factory.h"
+#include "shared/test/unit_test/helpers/debug_manager_state_restore.h"
 #include "shared/test/unit_test/helpers/default_hw_info.h"
 #include "shared/test/unit_test/mocks/mock_device.h"
 
@@ -33,5 +35,23 @@ struct DeviceFixture {
     L0::Device *device = nullptr;
 };
 
+struct MultiDeviceFixture {
+    virtual void SetUp() { // NOLINT(readability-identifier-naming)
+        DebugManager.flags.CreateMultipleRootDevices.set(numRootDevices);
+        auto executionEnvironment = new NEO::ExecutionEnvironment;
+        auto devices = NEO::DeviceFactory::createDevices(*executionEnvironment);
+        driverHandle = std::make_unique<Mock<L0::DriverHandleImp>>();
+        driverHandle->initialize(std::move(devices));
+    }
+
+    virtual void TearDown() { // NOLINT(readability-identifier-naming)
+    }
+
+    DebugManagerStateRestore restorer;
+    std::unique_ptr<Mock<L0::DriverHandleImp>> driverHandle;
+    std::vector<L0::Device *> devices;
+    const uint32_t numRootDevices = 4u;
+};
+
 } // namespace ult
 } // namespace L0

level_zero/core/test/unit_tests/fixtures/module_fixture.h

@@ -65,5 +65,44 @@ struct ModuleFixture : public DeviceFixture {
     std::unique_ptr<WhiteBox<::L0::Kernel>> kernel;
 };
 
+struct MultiDeviceModuleFixture : public MultiDeviceFixture {
+    void SetUp() override {
+        MultiDeviceFixture::SetUp();
+        modules.resize(numRootDevices);
+    }
+
+    void createModuleFromBinary(uint32_t rootDeviceIndex) {
+        std::string testFile;
+        retrieveBinaryKernelFilename(testFile, binaryFilename + "_", ".bin");
+
+        size_t size = 0;
+        auto src = loadDataFromFile(testFile.c_str(), size);
+
+        ASSERT_NE(0u, size);
+        ASSERT_NE(nullptr, src);
+
+        ze_module_desc_t moduleDesc = {ZE_MODULE_DESC_VERSION_CURRENT};
+        moduleDesc.format = ZE_MODULE_FORMAT_NATIVE;
+        moduleDesc.pInputModule = reinterpret_cast<const uint8_t *>(src.get());
+        moduleDesc.inputSize = size;
+
+        ModuleBuildLog *moduleBuildLog = nullptr;
+
+        auto device = driverHandle->devices[rootDeviceIndex];
+        modules[rootDeviceIndex].reset(Module::create(device,
+                                                      &moduleDesc,
+                                                      moduleBuildLog));
+    }
+
+    void TearDown() override {
+        MultiDeviceFixture::TearDown();
+    }
+
+    const std::string binaryFilename = "test_kernel";
+    const std::string kernelName = "test";
+    const uint32_t numKernelArguments = 6;
+    std::vector<std::unique_ptr<L0::Module>> modules;
+};
+
 } // namespace ult
 } // namespace L0

level_zero/core/test/unit_tests/sources/module/test_module.cpp

@@ -236,5 +236,52 @@ HWTEST_F(ModuleDynamicLinkTests, givenCallToDynamicLinkThenUnsupportedFeatureIsR
     delete module1;
 }
 
+class MultiDeviceModuleSetArgBufferTest : public MultiDeviceModuleFixture, public ::testing::Test {
+  public:
+    void SetUp() override {
+        MultiDeviceModuleFixture::SetUp();
+    }
+
+    void TearDown() override {
+        MultiDeviceModuleFixture::TearDown();
+    }
+
+    void createKernelAndAllocMemory(uint32_t rootDeviceIndex, void **ptr, ze_kernel_handle_t *kernelHandle) {
+        ze_kernel_desc_t kernelDesc = {};
+        kernelDesc.version = ZE_KERNEL_DESC_VERSION_CURRENT;
+        kernelDesc.flags = ZE_KERNEL_FLAG_NONE;
+        kernelDesc.pKernelName = kernelName.c_str();
+        ze_result_t res = modules[rootDeviceIndex].get()->createKernel(&kernelDesc, kernelHandle);
+        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+        res = driverHandle->allocHostMem(ZE_HOST_MEM_ALLOC_FLAG_DEFAULT, 4096u, rootDeviceIndex, ptr);
+        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    }
+};
+
+HWTEST_F(MultiDeviceModuleSetArgBufferTest,
+         givenCallsToSetArgBufferThenAllocationIsSetForCorrectDevice) {
+
+    for (uint32_t rootDeviceIndex = 0; rootDeviceIndex < numRootDevices; rootDeviceIndex++) {
+        createModuleFromBinary(rootDeviceIndex);
+
+        ze_kernel_handle_t kernelHandle;
+        void *ptr = nullptr;
+        createKernelAndAllocMemory(rootDeviceIndex, &ptr, &kernelHandle);
+
+        L0::KernelImp *kernel = reinterpret_cast<L0::KernelImp *>(Kernel::fromHandle(kernelHandle));
+        kernel->setArgBuffer(0, sizeof(ptr), &ptr);
+
+        for (auto alloc : kernel->getResidencyContainer()) {
+            if (alloc && alloc->getGpuAddress() == reinterpret_cast<uint64_t>(ptr)) {
+                EXPECT_EQ(rootDeviceIndex, alloc->getRootDeviceIndex());
+            }
+        }
+
+        driverHandle->freeMem(ptr);
+        Kernel::fromHandle(kernelHandle)->destroy();
+    }
+}
+
 } // namespace ult
 } // namespace L0