Optimize loads/stores on nested structs

Also fix LocalAccessChainConvert test: nested structs now convert

Add InsertExtractElim test for nested struct

3 files changed, 52 insertions, 26 deletions

diff --git a/source/opt/insert_extract_elim.cpp b/source/opt/insert_extract_elim.cpp
index d3c64854..b506d9eb 100644
--- a/source/opt/insert_extract_elim.cpp
+++ b/source/opt/insert_extract_elim.cpp
@@ -31,24 +31,28 @@ const uint32_t kInsertCompositeIdInIdx = 1;
 }  // anonymous namespace
 
 bool InsertExtractElimPass::ExtInsMatch(const ir::Instruction* extInst,
-                                        const ir::Instruction* insInst) const {
-  if (extInst->NumInOperands() != insInst->NumInOperands() - 1) return false;
-  uint32_t numIdx = extInst->NumInOperands() - 1;
+                                        const ir::Instruction* insInst,
+                                        const uint32_t extOffset) const {
+  if (extInst->NumInOperands() - extOffset != insInst->NumInOperands() - 1)
+    return false;
+  uint32_t numIdx = extInst->NumInOperands() - 1 - extOffset;
   for (uint32_t i = 0; i < numIdx; ++i)
-    if (extInst->GetSingleWordInOperand(i + 1) !=
+    if (extInst->GetSingleWordInOperand(i + 1 + extOffset) !=
         insInst->GetSingleWordInOperand(i + 2))
       return false;
   return true;
 }
 
 bool InsertExtractElimPass::ExtInsConflict(
-    const ir::Instruction* extInst, const ir::Instruction* insInst) const {
-  if (extInst->NumInOperands() == insInst->NumInOperands() - 1) return false;
-  uint32_t extNumIdx = extInst->NumInOperands() - 1;
+    const ir::Instruction* extInst, const ir::Instruction* insInst,
+    const uint32_t extOffset) const {
+  if (extInst->NumInOperands() - extOffset == insInst->NumInOperands() - 1)
+    return false;
+  uint32_t extNumIdx = extInst->NumInOperands() - 1 - extOffset;
   uint32_t insNumIdx = insInst->NumInOperands() - 2;
   uint32_t numIdx = std::min(extNumIdx, insNumIdx);
   for (uint32_t i = 0; i < numIdx; ++i)
-    if (extInst->GetSingleWordInOperand(i + 1) !=
+    if (extInst->GetSingleWordInOperand(i + 1 + extOffset) !=
         insInst->GetSingleWordInOperand(i + 2))
       return false;
   return true;
@@ -68,13 +72,32 @@ bool InsertExtractElimPass::EliminateInsertExtract(ir::Function* func) {
           uint32_t cid = ii->GetSingleWordInOperand(kExtractCompositeIdInIdx);
           ir::Instruction* cinst = get_def_use_mgr()->GetDef(cid);
           uint32_t replId = 0;
+          // Offset of extract indices being compared to insert indices.
+          // Offset increases as indices are matched.
+          uint32_t extOffset = 0;
           while (cinst->opcode() == SpvOpCompositeInsert) {
-            if (ExtInsConflict(&*ii, cinst)) break;
-            if (ExtInsMatch(&*ii, cinst)) {
+            if (ExtInsMatch(&*ii, cinst, extOffset)) {
+              // Match! Use inserted value as replacement
               replId = cinst->GetSingleWordInOperand(kInsertObjectIdInIdx);
               break;
             }
-            cid = cinst->GetSingleWordInOperand(kInsertCompositeIdInIdx);
+            else if (ExtInsConflict(&*ii, cinst, extOffset)) {
+              // If extract has fewer indices than the insert, stop searching.
+              // Otherwise increment offset of extract indices considered and
+              // continue searching through the inserted value
+              if (ii->NumInOperands() - extOffset <
+                  cinst->NumInOperands() - 1) {
+                break;
+              }
+              else {
+                extOffset += cinst->NumInOperands() - 2;
+                cid = cinst->GetSingleWordInOperand(kInsertObjectIdInIdx);
+              }
+            }
+            else {
+              // Consider next composite in insert chain
+              cid = cinst->GetSingleWordInOperand(kInsertCompositeIdInIdx);
+            }
             cinst = get_def_use_mgr()->GetDef(cid);
           }
           // If search ended with CompositeConstruct or ConstantComposite
@@ -85,8 +108,8 @@ bool InsertExtractElimPass::EliminateInsertExtract(ir::Function* func) {
           // vector composition, and additional CompositeInsert.
           if ((cinst->opcode() == SpvOpCompositeConstruct ||
                cinst->opcode() == SpvOpConstantComposite) &&
-              (*ii).NumInOperands() == 2) {
-            uint32_t compIdx = (*ii).GetSingleWordInOperand(1);
+              (*ii).NumInOperands() - extOffset == 2) {
+            uint32_t compIdx = (*ii).GetSingleWordInOperand(extOffset + 1);
             if (IsVectorType(cinst->type_id())) {
               if (compIdx < cinst->NumInOperands()) {
                 uint32_t i = 0;
diff --git a/source/opt/insert_extract_elim.h b/source/opt/insert_extract_elim.h
index d5dba00f..cb0f4f3f 100644
--- a/source/opt/insert_extract_elim.h
+++ b/source/opt/insert_extract_elim.h
@@ -40,25 +40,28 @@ class InsertExtractElimPass : public Pass {
   Status Process(ir::IRContext*) override;
 
  private:
-  // Return true if indices of extract |extInst| and insert |insInst| match
+  // Return true if indices of extract |extInst| starting at |extOffset|
+  // match indices of insert |insInst|. 
   bool ExtInsMatch(const ir::Instruction* extInst,
-                   const ir::Instruction* insInst) const;
+                   const ir::Instruction* insInst,
+                   const uint32_t extOffset) const;
 
-  // Return true if indices of extract |extInst| and insert |insInst| conflict,
-  // specifically, if the insert changes bits specified by the extract, but
-  // changes either more bits or less bits than the extract specifies,
-  // meaning the exact value being inserted cannot be used to replace
-  // the extract.
+  // Return true if indices of extract |extInst| starting at |extOffset| and
+  // indices of insert |insInst| conflict, specifically, if the insert
+  // changes bits specified by the extract, but changes either more bits
+  // or less bits than the extract specifies, meaning the exact value being
+  // inserted cannot be used to replace the extract.
   bool ExtInsConflict(const ir::Instruction* extInst,
-                      const ir::Instruction* insInst) const;
+                      const ir::Instruction* insInst,
+                      const uint32_t extOffset) const;
 
   // Return true if |typeId| is a vector type
   bool IsVectorType(uint32_t typeId);
 
-  // Look for OpExtract on sequence of OpInserts in |func|. If there is an
-  // insert with identical indices, replace the extract with the value
-  // that is inserted if possible. Specifically, replace if there is no
-  // intervening insert which conflicts.
+  // Look for OpExtract on sequence of OpInserts in |func|. If there is a
+  // reaching insert which corresponds to the indices of the extract, replace
+  // the extract with the value that is inserted. Also resolve extracts from
+  // CompositeConstruct or ConstantComposite.
   bool EliminateInsertExtract(ir::Function* func);
 
   // Initialize extensions whitelist
diff --git a/source/opt/mem_pass.cpp b/source/opt/mem_pass.cpp
index ae5baa81..1292804a 100644
--- a/source/opt/mem_pass.cpp
+++ b/source/opt/mem_pass.cpp
@@ -64,7 +64,7 @@ bool MemPass::IsTargetType(const ir::Instruction* typeInst) const {
   int nonMathComp = 0;
   typeInst->ForEachInId([&nonMathComp, this](const uint32_t* tid) {
     ir::Instruction* compTypeInst = get_def_use_mgr()->GetDef(*tid);
-    if (!IsBaseTargetType(compTypeInst)) ++nonMathComp;
+    if (!IsTargetType(compTypeInst)) ++nonMathComp;
   });
   return nonMathComp == 0;
 }