refactor MemPool options field

Author: adurang

offload/plugins-nextgen/level_zero/include/L0Options.h

@@ -77,13 +77,18 @@ struct L0OptionsTy {
   /// Staging buffer count
   size_t StagingBufferCount = L0StagingBufferCount;
 
-  // TODO: This should probably be an array indexed by AllocKind
   /// Memory pool parameters
   /// MemPoolInfo[MemType] = {AllocMax(MB), Capacity, PoolSize(MB)}
-  std::map<int32_t, std::array<int32_t, 3>> MemPoolInfo = {
-      {TARGET_ALLOC_DEVICE, {1, 4, 256}},
-      {TARGET_ALLOC_HOST, {1, 4, 256}},
-      {TARGET_ALLOC_SHARED, {8, 4, 256}}};
+  struct MemPoolConfigTy {
+    bool Use;
+    int32_t AllocMax;
+    int32_t Capacity;
+    int32_t PoolSize;
+  };
+  std::array<MemPoolConfigTy, 3> MemPoolConfig{
+      MemPoolConfigTy{true, 1, 4, 256},  // TARGET_ALLOC_DEVICE
+      MemPoolConfigTy{true, 1, 4, 256},  // TARGET_ALLOC_HOST
+      MemPoolConfigTy{true, 8, 4, 256}}; // TARGET_ALLOC_SHARED
 
   /// Parameters for memory pools dedicated to reduction scratch space
   std::array<int32_t, 3> ReductionPoolInfo{256, 8, 8192};

offload/plugins-nextgen/level_zero/src/L0Memory.cpp

@@ -69,10 +69,10 @@ Error MemAllocatorTy::MemPoolTy::init(int32_t Kind, MemAllocatorTy *AllocatorIn,
   Allocator = AllocatorIn;
 
   // Read user-defined options
-  const auto &UserOptions = Option.MemPoolInfo.at(AllocKind);
-  const size_t UserAllocMax = UserOptions[0];
-  const size_t UserCapacity = UserOptions[1];
-  const size_t UserPoolSize = UserOptions[2];
+  const auto &UserOptions = Option.MemPoolConfig[AllocKind];
+  const size_t UserAllocMax = UserOptions.AllocMax;
+  const size_t UserCapacity = UserOptions.Capacity;
+  const size_t UserPoolSize = UserOptions.PoolSize;
 
   BlockCapacity = UserCapacity;
   PoolSizeMax = UserPoolSize << 20; // MB to B
@@ -373,7 +373,7 @@ Error MemAllocatorTy::initDevicePools(L0DeviceTy &L0Device,
   Device = &L0Device;
   L0Context = &L0Device.getL0Context();
   for (auto Kind : {TARGET_ALLOC_DEVICE, TARGET_ALLOC_SHARED}) {
-    if (Options.MemPoolInfo.count(Kind) > 0) {
+    if (Options.MemPoolConfig[Kind].Use) {
       std::lock_guard<std::mutex> Lock(Mtx);
       Pools[Kind] = std::make_unique<MemPoolTy>();
       if (auto Err = Pools[Kind]->init(Kind, this, Options))
@@ -395,7 +395,7 @@ Error MemAllocatorTy::initHostPool(L0ContextTy &Driver,
   SupportsLargeMem = Driver.supportsLargeMem();
   IsHostMem = true;
   this->L0Context = &Driver;
-  if (Option.MemPoolInfo.count(TARGET_ALLOC_HOST) > 0) {
+  if (Option.MemPoolConfig[TARGET_ALLOC_HOST].Use) {
     std::lock_guard<std::mutex> Lock(Mtx);
     Pools[TARGET_ALLOC_HOST] = std::make_unique<MemPoolTy>();
     if (auto Err =

offload/plugins-nextgen/level_zero/src/L0Options.cpp

@@ -41,7 +41,8 @@ void L0OptionsTy::processEnvironmentVars() {
   if (MemoryPoolVar.isPresent()) {
     if (MemoryPoolVar.get() == "0") {
       Flags.UseMemoryPool = 0;
-      MemPoolInfo.clear();
+      std::for_each(MemPoolConfig.begin(), MemPoolConfig.end(),
+                    [](auto &I) { I = {false, 0, 0, 0}; });
     } else {
       std::istringstream Str(MemoryPoolVar.get());
       int32_t MemType = -1;
@@ -50,19 +51,19 @@ void L0OptionsTy::processEnvironmentVars() {
       const std::array<int32_t, 3> DefaultValue{1, 4, 256};
       const int32_t AllMemType = INT32_MAX;
       std::array<int32_t, 3> AllInfo{1, 4, 256};
-      std::map<int32_t, std::array<int32_t, 3>> PoolInfo;
+      std::array<std::array<int32_t, 3>, 3> PoolInfo;
       for (std::string Token; std::getline(Str, Token, ',') && Valid > 0;) {
         if (Token == "device") {
           MemType = TARGET_ALLOC_DEVICE;
-          PoolInfo.emplace(MemType, DefaultValue);
+          PoolInfo[TARGET_ALLOC_DEVICE] = DefaultValue;
           Offset = 0;
         } else if (Token == "host") {
           MemType = TARGET_ALLOC_HOST;
-          PoolInfo.emplace(MemType, DefaultValue);
+          PoolInfo[TARGET_ALLOC_HOST] = DefaultValue;
           Offset = 0;
         } else if (Token == "shared") {
           MemType = TARGET_ALLOC_SHARED;
-          PoolInfo.emplace(MemType, DefaultValue);
+          PoolInfo[TARGET_ALLOC_SHARED] = DefaultValue;
           Offset = 0;
         } else if (Token == "all") {
           MemType = AllMemType;
@@ -87,19 +88,24 @@ void L0OptionsTy::processEnvironmentVars() {
         if (Valid == 2) {
           // "all" is specified -- ignore other inputs
           if (AllInfo[0] > 0) {
-            MemPoolInfo[TARGET_ALLOC_DEVICE] = AllInfo;
-            MemPoolInfo[TARGET_ALLOC_HOST] = AllInfo;
-            MemPoolInfo[TARGET_ALLOC_SHARED] = std::move(AllInfo);
+            MemPoolConfig[TARGET_ALLOC_DEVICE] = {true, AllInfo[0], AllInfo[1],
+                                                  AllInfo[2]};
+            MemPoolConfig[TARGET_ALLOC_HOST] = {true, AllInfo[0], AllInfo[1],
+                                                AllInfo[2]};
+            MemPoolConfig[TARGET_ALLOC_SHARED] = {true, AllInfo[0], AllInfo[1],
+                                                  AllInfo[2]};
           } else {
-            MemPoolInfo.clear();
+            std::for_each(MemPoolConfig.begin(), MemPoolConfig.end(),
+                          [](auto &I) { I = {false, 0, 0, 0}; });
           }
         } else {
-          // Use user-specified configuration
-          for (auto &I : PoolInfo) {
-            if (I.second[0] > 0)
-              MemPoolInfo[I.first] = I.second;
-            else
-              MemPoolInfo.erase(I.first);
+          for (size_t Pool = 0; Pool < PoolInfo.size(); ++Pool) {
+            if (PoolInfo[Pool][0] == 0) {
+              MemPoolConfig[Pool] = {false, 0, 0, 0};
+            } else {
+              MemPoolConfig[Pool] = {true, PoolInfo[Pool][0], PoolInfo[Pool][1],
+                                     PoolInfo[Pool][2]};
+            }
           }
         }
       } else {