[llvm] [Matrix] Propagate shape information through Select insts (PR #141876)

[Florian Hahn via llvm-commits]

================
@@ -269,6 +270,16 @@ computeShapeInfoForInst(Instruction *I,
       return OpShape->second;
   }
 
+  if (auto *Select = dyn_cast<SelectInst>(I)) {
+    Type *CondTy = Select->getCondition()->getType();
+    for (Use &Op : CondTy->isVectorTy() ? Select->operands()
+                                        : drop_begin(Select->operands())) {
+      auto OpShape = ShapeMap.find(Op);
+      if (OpShape != ShapeMap.end())
+        return OpShape->second;
+    }
+  }
----------------
fhahn wrote:

```suggestion
  if (isUniformShape(I) || isa<SelectInst>(I) {
    auto ShapedOps = isUniformShape(I) ? I->operands() : drop_begin(I->operands());
    // Find the first operand that has a known shape and use that.
```

Not sure if we should include the shape of the condition if it is a vector; as the result will be a mix of the last 2 ops. If we skip the condition, we could merge it with isUniformShape below.

https://github.com/llvm/llvm-project/pull/141876