feature: add support for zeKernelGetAllocationPropertiesExp L0 extension

Related-To: GSD-9624

Signed-off-by: Igor Venevtsev <igor.venevtsev@intel.com>

level_zero/api/core/ze_module_api_entrypoints.h

@@ -270,7 +270,7 @@ ze_result_t ZE_APICALL zeKernelGetAllocationPropertiesExp(
     ze_kernel_handle_t hKernel,
     uint32_t *pCount,
     ze_kernel_allocation_exp_properties_t *pAllocationProperties) {
-    return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    return L0::Kernel::fromHandle(hKernel)->getAllocationProperties(pCount, pAllocationProperties);
 }
 } // namespace L0
 

level_zero/core/source/driver/driver_handle_imp_helper.cpp

@@ -62,6 +62,7 @@ const std::vector<std::pair<std::string, uint32_t>> DriverHandleImp::extensionsS
     {ZEX_INTEL_QUEUE_COPY_OPERATIONS_OFFLOAD_HINT_EXP_NAME, ZEX_INTEL_QUEUE_COPY_OPERATIONS_OFFLOAD_HINT_EXP_VERSION_CURRENT},
     {ZE_FABRIC_EXP_NAME, ZE_FABRIC_EXP_VERSION_CURRENT},
     {ZE_BANDWIDTH_PROPERTIES_EXP_NAME, ZE_BANDWIDTH_PROPERTIES_EXP_VERSION_CURRENT},
+    {ZE_GET_KERNEL_ALLOCATION_PROPERTIES_EXP_NAME, ZE_KERNEL_GET_ALLOCATION_PROPERTIES_EXP_VERSION_CURRENT},
 
     // Metrics Driver experimental extensions
     {ZET_INTEL_METRIC_APPEND_MARKER_EXP_NAME, ZET_INTEL_METRIC_APPEND_MARKER_EXP_VERSION_CURRENT},

level_zero/core/source/kernel/kernel.h

@@ -125,6 +125,7 @@ struct Kernel : _ze_kernel_handle_t, virtual NEO::DispatchKernelEncoderI, NEO::N
     virtual ze_result_t destroy() = 0;
     virtual ze_result_t getBaseAddress(uint64_t *baseAddress) = 0;
     virtual ze_result_t getKernelProgramBinary(size_t *kernelSize, char *pKernelBinary) = 0;
+    virtual ze_result_t getAllocationProperties(uint32_t *pCount, ze_kernel_allocation_exp_properties_t *pAllocationProperties) = 0;
     virtual ze_result_t setIndirectAccess(ze_kernel_indirect_access_flags_t flags) = 0;
     virtual ze_result_t getIndirectAccess(ze_kernel_indirect_access_flags_t *flags) = 0;
     virtual ze_result_t getSourceAttributes(uint32_t *pSize, char **pString) = 0;

level_zero/core/source/kernel/kernel_imp.cpp

@@ -379,6 +379,56 @@ ze_result_t KernelImp::getKernelProgramBinary(size_t *kernelSize, char *pKernelB
     return ZE_RESULT_SUCCESS;
 }
 
+ze_result_t KernelImp::getAllocationProperties(uint32_t *pCount, ze_kernel_allocation_exp_properties_t *pAllocationProperties) {
+    uint32_t totalAllocations = 0;
+
+    for (const auto &alloc : getArgumentsResidencyContainer()) {
+        if (alloc) {
+            ++totalAllocations;
+        }
+    }
+
+    totalAllocations += static_cast<uint32_t>(getInternalResidencyContainer().size());
+
+    if (*pCount == 0 || *pCount > totalAllocations) {
+        *pCount = totalAllocations;
+    }
+
+    // If pAllocationProperties is nullptr, then user getting *pCount first and calling second time
+    if (pAllocationProperties == nullptr) {
+        return ZE_RESULT_SUCCESS;
+    }
+
+    auto svmAllocsManager = this->module->getDevice()->getDriverHandle()->getSvmAllocsManager();
+    uint32_t allocIndex = 0;
+
+    auto parseResidencyContainer = [&](const NEO::ResidencyContainer &residencyContainer, bool kernelArgumentsContainer) {
+        for (uint32_t i = 0; allocIndex < *pCount && i < residencyContainer.size(); ++i) {
+            if (!residencyContainer[i]) {
+                continue;
+            }
+
+            pAllocationProperties[allocIndex].base = residencyContainer[i]->getGpuAddress();
+            if (auto svmAlloc = svmAllocsManager->getSVMAlloc(reinterpret_cast<const void *>(residencyContainer[i]->getGpuAddress()))) {
+                pAllocationProperties[allocIndex].size = svmAlloc->size;
+                pAllocationProperties[allocIndex].type = Context::parseUSMType(svmAlloc->memoryType);
+            } else {
+                pAllocationProperties[allocIndex].size = residencyContainer[i]->getUnderlyingBufferSize();
+                pAllocationProperties[allocIndex].type = ZE_MEMORY_TYPE_UNKNOWN;
+            }
+
+            pAllocationProperties[allocIndex].argIndex = kernelArgumentsContainer ? i : std::numeric_limits<uint32_t>::max();
+
+            ++allocIndex;
+        }
+    };
+
+    parseResidencyContainer(getArgumentsResidencyContainer(), true);
+    parseResidencyContainer(getInternalResidencyContainer(), false);
+
+    return ZE_RESULT_SUCCESS;
+}
+
 ze_result_t KernelImp::setArgumentValue(uint32_t argIndex, size_t argSize,
                                         const void *pArgValue) {
     if (argIndex >= privateState.kernelArgHandlers.size()) {

level_zero/core/source/kernel/kernel_imp.h

@@ -47,6 +47,7 @@ struct KernelImp : Kernel {
 
     ze_result_t getBaseAddress(uint64_t *baseAddress) override;
     ze_result_t getKernelProgramBinary(size_t *kernelSize, char *pKernelBinary) override;
+    ze_result_t getAllocationProperties(uint32_t *pCount, ze_kernel_allocation_exp_properties_t *pAllocationProperties) override;
     ze_result_t setIndirectAccess(ze_kernel_indirect_access_flags_t flags) override;
     ze_result_t getIndirectAccess(ze_kernel_indirect_access_flags_t *flags) override;
     ze_result_t getSourceAttributes(uint32_t *pSize, char **pString) override;

level_zero/core/test/unit_tests/sources/driver/test_driver.cpp

@@ -1985,6 +1985,7 @@ TEST_F(DriverExtensionsTest, givenDriverHandleWhenAskingForExtensionsThenReturnC
     verifyExtensionDefinition(ZEX_INTEL_QUEUE_COPY_OPERATIONS_OFFLOAD_HINT_EXP_NAME, ZEX_INTEL_QUEUE_COPY_OPERATIONS_OFFLOAD_HINT_EXP_VERSION_CURRENT);
     verifyExtensionDefinition(ZE_FABRIC_EXP_NAME, ZE_FABRIC_EXP_VERSION_CURRENT);
     verifyExtensionDefinition(ZE_BANDWIDTH_PROPERTIES_EXP_NAME, ZE_BANDWIDTH_PROPERTIES_EXP_VERSION_CURRENT);
+    verifyExtensionDefinition(ZE_GET_KERNEL_ALLOCATION_PROPERTIES_EXP_NAME, ZE_KERNEL_GET_ALLOCATION_PROPERTIES_EXP_VERSION_CURRENT);
 }
 
 } // namespace ult

level_zero/core/test/unit_tests/sources/kernel/test_kernel.cpp

@@ -1810,7 +1810,7 @@ TEST_F(KernelPropertiesTests, givenValidKernelThenPropertiesAreRetrieved) {
     ze_result_t res = kernel->getProperties(&kernelProperties);
     EXPECT_EQ(ZE_RESULT_SUCCESS, res);
 
-    EXPECT_EQ(6U, kernelProperties.numKernelArgs);
+    EXPECT_EQ(7U, kernelProperties.numKernelArgs);
 
     EXPECT_EQ(0U, kernelProperties.requiredNumSubGroups);
     EXPECT_EQ(0U, kernelProperties.requiredSubgroupSize);

level_zero/core/test/unit_tests/sources/kernel/test_kernel_2.cpp

@@ -21,6 +21,7 @@
 #include "shared/test/common/test_macros/hw_test.h"
 #include "shared/test/common/test_macros/test.h"
 
+#include "level_zero/core/source/context/context_imp.h"
 #include "level_zero/core/source/kernel/kernel_shared_state.h"
 #include "level_zero/core/test/unit_tests/fixtures/device_fixture.h"
 #include "level_zero/core/test/unit_tests/fixtures/kernel_max_cooperative_groups_count_fixture.h"
@@ -1388,5 +1389,289 @@ TEST_F(KernelImpTest, givenDefaultGroupSizeWhenGetGroupSizeCalledThenReturnDefau
     EXPECT_EQ(1u, groupSize[2]);
 }
 
+struct KernelAllocationPropertiesExpFixture : ModuleFixture {
+    void createSimpleKernel() {
+        kernel = createKernelWithName("memcpy_bytes_attr"); // simple kernel w/o printf, no internal allocs to report
+        hKernel = kernel->toHandle();
+        hDevice = device->toHandle();
+        hContext = context->toHandle();
+    }
+
+    void createKernelWithPrintf() {
+        kernel = createKernelWithName("test"); // test kernel has printf, i.e one internal alloc to report
+        hKernel = kernel->toHandle();
+        hDevice = device->toHandle();
+        hContext = context->toHandle();
+    }
+
+    ze_kernel_handle_t hKernel = nullptr;
+    ze_context_handle_t hContext = nullptr;
+    ze_device_handle_t hDevice = nullptr;
+
+    ze_device_mem_alloc_desc_t deviceMemDesc = {
+        .stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC,
+        .pNext = nullptr,
+        .flags = ZE_DEVICE_MEM_ALLOC_FLAG_BIAS_UNCACHED,
+        .ordinal = 0};
+
+    ze_host_mem_alloc_desc_t hostMemDesc = {
+        .stype = ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC,
+        .pNext = nullptr,
+        .flags = ZE_HOST_MEM_ALLOC_FLAG_BIAS_UNCACHED};
+};
+
+using KernelAllocationPropertiesExpTest = Test<KernelAllocationPropertiesExpFixture>;
+
+TEST_F(KernelAllocationPropertiesExpTest, givenSimpleKernelAndNoArgsSetThenNoAllocationPropertiesReturned) {
+    createSimpleKernel();
+
+    uint32_t count = 0;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, nullptr));
+    EXPECT_EQ(0u, count);
+}
+
+TEST_F(KernelAllocationPropertiesExpTest, givenSimpleKernelAndArgsSetThenCorrectAllocationPropertiesReturned) {
+    createSimpleKernel();
+
+    void *hostPtr = nullptr;
+    size_t hostPtrSize = 1024u;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemAllocHost(hContext, &hostMemDesc, hostPtrSize, 1, &hostPtr));
+
+    void *devicePtr = nullptr;
+    size_t devicePtrSize = 2048u;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemAllocDevice(hContext, &deviceMemDesc, devicePtrSize, 1, hDevice, &devicePtr));
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelSetArgumentValue(hKernel, 0, sizeof(hostPtr), &hostPtr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelSetArgumentValue(hKernel, 1, sizeof(devicePtr), &devicePtr));
+
+    uint32_t count = 0;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, nullptr));
+    EXPECT_EQ(2u, count);
+
+    std::vector<ze_kernel_allocation_exp_properties_t> kernelAllocationProps(count);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, kernelAllocationProps.data()));
+    EXPECT_EQ(2u, count);
+
+    EXPECT_EQ(reinterpret_cast<uint64_t>(hostPtr), kernelAllocationProps[0].base);
+    EXPECT_EQ(hostPtrSize, kernelAllocationProps[0].size);
+    EXPECT_EQ(ZE_MEMORY_TYPE_HOST, kernelAllocationProps[0].type);
+    EXPECT_EQ(0u, kernelAllocationProps[0].argIndex);
+
+    EXPECT_EQ(reinterpret_cast<uint64_t>(devicePtr), kernelAllocationProps[1].base);
+    EXPECT_EQ(devicePtrSize, kernelAllocationProps[1].size);
+    EXPECT_EQ(ZE_MEMORY_TYPE_DEVICE, kernelAllocationProps[1].type);
+    EXPECT_EQ(1u, kernelAllocationProps[1].argIndex);
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemFree(hContext, hostPtr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemFree(hContext, devicePtr));
+}
+
+TEST_F(KernelAllocationPropertiesExpTest, givenKernelWithInternalAllocationAndNoArgsSetThenCorrectAllocationPropertiesReturned) {
+    createKernelWithPrintf();
+
+    uint32_t count = 0;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, nullptr));
+    EXPECT_EQ(1u /* printf */, count);
+
+    std::vector<ze_kernel_allocation_exp_properties_t> kernelAllocationProps(count);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, kernelAllocationProps.data()));
+    EXPECT_EQ(1u, count);
+
+    auto printfBuffer = kernel->sharedState->printfBuffer;
+    EXPECT_EQ(printfBuffer->getGpuAddress(), kernelAllocationProps[0].base);
+    EXPECT_EQ(printfBuffer->getUnderlyingBufferSize(), kernelAllocationProps[0].size);
+    EXPECT_EQ(ZE_MEMORY_TYPE_UNKNOWN, kernelAllocationProps[0].type);
+    EXPECT_EQ(std::numeric_limits<uint32_t>::max(), kernelAllocationProps[0].argIndex);
+}
+
+TEST_F(KernelAllocationPropertiesExpTest, givenKernelWithInternalAllocationAndArgsSetThenCorrectAllocationPropertiesReturned) {
+    createKernelWithPrintf();
+
+    void *hostPtr = nullptr;
+    size_t hostPtrSize = 1024u;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemAllocHost(hContext, &hostMemDesc, hostPtrSize, 1, &hostPtr));
+
+    void *devicePtr = nullptr;
+    size_t devicePtrSize = 2048u;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemAllocDevice(hContext, &deviceMemDesc, devicePtrSize, 1, hDevice, &devicePtr));
+
+    void *sharedPtr = nullptr;
+    size_t sharedPtrSize = 4096u;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemAllocShared(hContext, &deviceMemDesc, &hostMemDesc, sharedPtrSize, 1, hDevice, &sharedPtr));
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelSetArgumentValue(hKernel, 0, sizeof(hostPtr), &hostPtr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelSetArgumentValue(hKernel, 1, sizeof(devicePtr), &devicePtr));
+    // No args 2-5 set
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelSetArgumentValue(hKernel, 6, sizeof(sharedPtr), &sharedPtr));
+
+    uint32_t count = 0;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, nullptr));
+    EXPECT_EQ(4u /* 3 args + printf */, count);
+
+    std::vector<ze_kernel_allocation_exp_properties_t> kernelAllocationProps(count);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, kernelAllocationProps.data()));
+    EXPECT_EQ(4u, count);
+
+    EXPECT_EQ(reinterpret_cast<uint64_t>(hostPtr), kernelAllocationProps[0].base);
+    EXPECT_EQ(hostPtrSize, kernelAllocationProps[0].size);
+    EXPECT_EQ(ZE_MEMORY_TYPE_HOST, kernelAllocationProps[0].type);
+    EXPECT_EQ(0u, kernelAllocationProps[0].argIndex);
+
+    EXPECT_EQ(reinterpret_cast<uint64_t>(devicePtr), kernelAllocationProps[1].base);
+    EXPECT_EQ(devicePtrSize, kernelAllocationProps[1].size);
+    EXPECT_EQ(ZE_MEMORY_TYPE_DEVICE, kernelAllocationProps[1].type);
+    EXPECT_EQ(1u, kernelAllocationProps[1].argIndex);
+
+    EXPECT_EQ(reinterpret_cast<uint64_t>(sharedPtr), kernelAllocationProps[2].base);
+    EXPECT_EQ(sharedPtrSize, kernelAllocationProps[2].size);
+    EXPECT_EQ(ZE_MEMORY_TYPE_SHARED, kernelAllocationProps[2].type);
+    EXPECT_EQ(6u, kernelAllocationProps[2].argIndex);
+
+    auto printfBuffer = kernel->sharedState->printfBuffer;
+    EXPECT_EQ(printfBuffer->getGpuAddress(), kernelAllocationProps[3].base);
+    EXPECT_EQ(printfBuffer->getUnderlyingBufferSize(), kernelAllocationProps[3].size);
+    EXPECT_EQ(ZE_MEMORY_TYPE_UNKNOWN, kernelAllocationProps[3].type);
+    EXPECT_EQ(std::numeric_limits<uint32_t>::max(), kernelAllocationProps[3].argIndex);
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemFree(hContext, hostPtr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemFree(hContext, devicePtr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemFree(hContext, sharedPtr));
+}
+
+TEST_F(KernelAllocationPropertiesExpTest, givenKernelWithInternalAllocationAndArgsSetAndCountSmallerOrTooBigThenCorrectNumberOfAllocationPropertiesReturned) {
+    createKernelWithPrintf();
+
+    void *hostPtr = nullptr;
+    size_t hostPtrSize = 1024u;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemAllocHost(hContext, &hostMemDesc, hostPtrSize, 1, &hostPtr));
+
+    void *devicePtr = nullptr;
+    size_t devicePtrSize = 2048u;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemAllocDevice(hContext, &deviceMemDesc, devicePtrSize, 1, hDevice, &devicePtr));
+
+    void *sharedPtr = nullptr;
+    size_t sharedPtrSize = 4096u;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemAllocShared(hContext, &deviceMemDesc, &hostMemDesc, sharedPtrSize, 1, hDevice, &sharedPtr));
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelSetArgumentValue(hKernel, 0, sizeof(hostPtr), &hostPtr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelSetArgumentValue(hKernel, 1, sizeof(devicePtr), &devicePtr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelSetArgumentValue(hKernel, 6, sizeof(sharedPtr), &sharedPtr));
+
+    uint32_t totalCount = 0;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &totalCount, nullptr));
+    EXPECT_EQ(4u /* 3 args + printf */, totalCount);
+
+    uint32_t count = 2;
+    std::vector<ze_kernel_allocation_exp_properties_t> kernelAllocationProps(100);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, kernelAllocationProps.data()));
+    EXPECT_EQ(2u, count);
+
+    count = 100;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, kernelAllocationProps.data()));
+    EXPECT_EQ(totalCount, count);
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemFree(hContext, hostPtr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemFree(hContext, devicePtr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemFree(hContext, sharedPtr));
+}
+
+TEST_F(KernelAllocationPropertiesExpTest, givenKernelWithAdditionalSectionsAndInternalAllocationAndNoArgsSetThenCorrectAllocationPropertiesReturned) {
+    auto elfAdditionalSections = {ZebinTestData::AppendElfAdditionalSection::global, ZebinTestData::AppendElfAdditionalSection::constant, ZebinTestData::AppendElfAdditionalSection::constantString};
+
+    zebinData = std::make_unique<ZebinTestData::ZebinWithL0TestCommonModule>(device->getHwInfo(), elfAdditionalSections);
+    const auto &src = zebinData->storage;
+
+    ze_module_desc_t moduleDesc = {};
+    moduleDesc.format = ZE_MODULE_FORMAT_NATIVE;
+    moduleDesc.pInputModule = reinterpret_cast<const uint8_t *>(src.data());
+    moduleDesc.inputSize = src.size();
+    module.reset(new WhiteBox<::L0::Module>{device, nullptr, ModuleType::user});
+    module->initialize(&moduleDesc, device->getNEODevice());
+
+    createKernelWithPrintf();
+
+    uint32_t count = 0;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, nullptr));
+    EXPECT_EQ(3u /* 0 args + globals + constants + printf */, count);
+
+    std::vector<ze_kernel_allocation_exp_properties_t> kernelAllocationProps(count);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, kernelAllocationProps.data()));
+    EXPECT_EQ(3u, count);
+
+    for (uint32_t i = 0; i < count; ++i) {
+        EXPECT_NE(0u, kernelAllocationProps[i].base);
+        EXPECT_NE(0u, kernelAllocationProps[i].size);
+        EXPECT_EQ(ZE_MEMORY_TYPE_UNKNOWN, kernelAllocationProps[i].type);
+        EXPECT_EQ(std::numeric_limits<uint32_t>::max(), kernelAllocationProps[i].argIndex);
+    }
+}
+
+TEST_F(KernelAllocationPropertiesExpTest, givenKernelWithAdditionalSectionsAndInternalAllocationAndArgsSetThenCorrectAllocationPropertiesReturned) {
+    auto elfAdditionalSections = {ZebinTestData::AppendElfAdditionalSection::global, ZebinTestData::AppendElfAdditionalSection::constant, ZebinTestData::AppendElfAdditionalSection::constantString};
+
+    zebinData = std::make_unique<ZebinTestData::ZebinWithL0TestCommonModule>(device->getHwInfo(), elfAdditionalSections);
+    const auto &src = zebinData->storage;
+
+    ze_module_desc_t moduleDesc = {};
+    moduleDesc.format = ZE_MODULE_FORMAT_NATIVE;
+    moduleDesc.pInputModule = reinterpret_cast<const uint8_t *>(src.data());
+    moduleDesc.inputSize = src.size();
+    module.reset(new WhiteBox<::L0::Module>{device, nullptr, ModuleType::user});
+    module->initialize(&moduleDesc, device->getNEODevice());
+
+    createKernelWithPrintf();
+
+    void *hostPtr = nullptr;
+    size_t hostPtrSize = 1024u;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemAllocHost(hContext, &hostMemDesc, hostPtrSize, 1, &hostPtr));
+
+    void *devicePtr = nullptr;
+    size_t devicePtrSize = 2048u;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemAllocDevice(hContext, &deviceMemDesc, devicePtrSize, 1, hDevice, &devicePtr));
+
+    void *sharedPtr = nullptr;
+    size_t sharedPtrSize = 4096u;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemAllocShared(hContext, &deviceMemDesc, &hostMemDesc, sharedPtrSize, 1, hDevice, &sharedPtr));
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelSetArgumentValue(hKernel, 0, sizeof(hostPtr), &hostPtr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelSetArgumentValue(hKernel, 1, sizeof(devicePtr), &devicePtr));
+    // No args 2-5 set
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelSetArgumentValue(hKernel, 6, sizeof(sharedPtr), &sharedPtr));
+
+    uint32_t count = 0;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, nullptr));
+    EXPECT_EQ(6u /* 3 args + globals + constants + printf */, count);
+
+    std::vector<ze_kernel_allocation_exp_properties_t> kernelAllocationProps(count);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeKernelGetAllocationPropertiesExp(hKernel, &count, kernelAllocationProps.data()));
+    EXPECT_EQ(6u, count);
+
+    EXPECT_EQ(reinterpret_cast<uint64_t>(hostPtr), kernelAllocationProps[0].base);
+    EXPECT_EQ(hostPtrSize, kernelAllocationProps[0].size);
+    EXPECT_EQ(ZE_MEMORY_TYPE_HOST, kernelAllocationProps[0].type);
+    EXPECT_EQ(0u, kernelAllocationProps[0].argIndex);
+
+    EXPECT_EQ(reinterpret_cast<uint64_t>(devicePtr), kernelAllocationProps[1].base);
+    EXPECT_EQ(devicePtrSize, kernelAllocationProps[1].size);
+    EXPECT_EQ(ZE_MEMORY_TYPE_DEVICE, kernelAllocationProps[1].type);
+    EXPECT_EQ(1u, kernelAllocationProps[1].argIndex);
+
+    EXPECT_EQ(reinterpret_cast<uint64_t>(sharedPtr), kernelAllocationProps[2].base);
+    EXPECT_EQ(sharedPtrSize, kernelAllocationProps[2].size);
+    EXPECT_EQ(ZE_MEMORY_TYPE_SHARED, kernelAllocationProps[2].type);
+    EXPECT_EQ(6u, kernelAllocationProps[2].argIndex);
+
+    for (uint32_t i = 3; i < count; ++i) {
+        EXPECT_NE(0u, kernelAllocationProps[i].base);
+        EXPECT_NE(0u, kernelAllocationProps[i].size);
+        EXPECT_EQ(ZE_MEMORY_TYPE_UNKNOWN, kernelAllocationProps[i].type);
+        EXPECT_EQ(std::numeric_limits<uint32_t>::max(), kernelAllocationProps[i].argIndex);
+    }
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemFree(hContext, hostPtr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemFree(hContext, devicePtr));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, zeMemFree(hContext, sharedPtr));
+}
+
 } // namespace ult
 } // namespace L0

shared/test/common/mocks/mock_modules_zebin.h

@@ -190,6 +190,20 @@ kernels:
         addrmode:        stateful
         addrspace:       image
         access_type:     writeonly
+      - arg_type:        arg_bypointer
+        offset:          0
+        size:            0
+        arg_index:       6
+        addrmode:        stateful
+        addrspace:       global
+        access_type:     readonly
+      - arg_type:        arg_bypointer
+        offset:          32
+        size:            8
+        arg_index:       6
+        addrmode:        stateless
+        addrspace:       global
+        access_type:     readonly
     per_thread_payload_arguments:
       - arg_type:        local_id
         offset:          0