[HdSt] Improve draw item batching in Hydra GL

Improve batching performance in Hydra GL, adding to PR #528. Thanks to Yoojin Jang for the initial idea.

The draw item "key" used in the map accounts for the geometric shader, buffer array versions and material params. Two draw items sharing the same "key" could use resources that don't aggregate. This checks happens in HdSt_DrawBatch::Append. So, a "key" effectively maps to several batches.

- Use a vector of batches as the value type in the map, instead of a single batch.
- Keep track of the previous draw item's key and batch. If the current draw item's key matches the last one, attempt to append it to the batch. If that fails, create a new batch and append it to the key's batch. If the key doesn't match, create a new entry into the map.

This should reduce the number of HdSt_DrawBatch's created to the necessary minimum, and improve GPU performance since more draw item's are bucketed into a batch than was the case earlier.

(Internal change: 1987160)

pxr/imaging/lib/hdSt/commandBuffer.cpp

@@ -46,6 +46,7 @@
 #include <tbb/enumerable_thread_specific.h>
 
 #include <functional>
+#include <unordered_map>
 
 PXR_NAMESPACE_OPEN_SCOPE
 
@@ -162,8 +163,27 @@ HdStCommandBuffer::_RebuildDrawBatches()
     bool bindlessTexture = GlfContextCaps::GetInstance()
                                                .bindlessTextureEnabled;
 
-    // XXX: Temporary sorting by shader.
-    std::map<size_t, HdSt_DrawBatchSharedPtr> batchMap;
+    // Use a cheap bucketing strategy to reduce to number of comparison tests
+    // required to figure out if a draw item can be batched.
+    // We use a hash of the geometric shader, BAR version and (optionally)
+    // material params as the key, and test (in the worst case) against each of 
+    // the batches for the key.
+    // Test against the previous draw item's hash and batch prior to looking up
+    // the map.
+    struct _PrevBatchHit {
+        _PrevBatchHit() : key(0) {}
+        void Update(size_t _key, HdSt_DrawBatchSharedPtr &_batch) {
+            key = _key;
+            batch = _batch;
+        }
+        size_t key;
+        HdSt_DrawBatchSharedPtr batch;
+    };
+    _PrevBatchHit prevBatch;
+
+    using _DrawBatchMap = 
+        std::unordered_map<size_t, HdSt_DrawBatchSharedPtrVector>;
+    _DrawBatchMap batchMap;
 
     for (size_t i = 0; i < _drawItems.size(); i++) {
         HdStDrawItem const * drawItem = _drawItems[i];
@@ -180,7 +200,6 @@ HdStCommandBuffer::_RebuildDrawBatches()
 
         size_t key = drawItem->GetGeometricShader()->ComputeHash();
         boost::hash_combine(key, drawItem->GetBufferArraysHash());
-
         if (!bindlessTexture) {
             // Geometric, RenderPass and Lighting shaders should never break
             // batches, however materials can. We consider the material 
@@ -192,14 +211,40 @@ HdStCommandBuffer::_RebuildDrawBatches()
         TF_DEBUG(HDST_DRAW_BATCH).Msg("%lu (%lu)\n", 
                 key, 
                 drawItem->GetBufferArraysHash());
-                //, drawItem->GetRprimID().GetText());
 
-        HdSt_DrawBatchSharedPtr batch;
-        TfMapLookup(batchMap, key, &batch);
-        if (!batch || !batch->Append(drawItemInstance)) {
-            batch = _NewDrawBatch(drawItemInstance);
-            _drawBatches.push_back(batch);
-            batchMap[key] = batch;
+        // Do a quick check to see if the draw item can be batched with the
+        // previous draw item, before checking the batchMap.
+        if (key == prevBatch.key && prevBatch.batch) {
+            if (prevBatch.batch->Append(drawItemInstance)) {
+                continue;
+            }
+        }
+
+        _DrawBatchMap::iterator const batchIter = batchMap.find(key);
+        bool const foundKey = batchIter != batchMap.end();
+        bool batched = false;
+        if (foundKey) {
+            HdSt_DrawBatchSharedPtrVector &batches = batchIter->second;
+            for (HdSt_DrawBatchSharedPtr &batch : batches) {
+                if (batch->Append(drawItemInstance)) {
+                    batched = true;
+                    prevBatch.Update(key, batch);
+                    break;
+                }
+            }
+        }
+
+        if (!batched) {
+            HdSt_DrawBatchSharedPtr batch = _NewDrawBatch(drawItemInstance);
+            _drawBatches.emplace_back(batch);
+            prevBatch.Update(key, batch);
+
+            if (foundKey) {
+                HdSt_DrawBatchSharedPtrVector &batches = batchIter->second;
+                batches.emplace_back(batch);
+            } else {
+                batchMap[key] = HdSt_DrawBatchSharedPtrVector({batch});
+            }
         }
     }
 }