[Mlir-commits] [mlir] [mlir][tosa]: Add SLICE_SHAPE folder (PR #186997)

[Luke Hutton]

================
@@ -2093,6 +2103,52 @@ OpFoldResult concatShapeFold(tosa::ConcatShapeOp *op) {
   return DenseElementsAttr::get(rankedTy, concatDims);
 }
 
+OpFoldResult sliceShapeFold(tosa::SliceShapeOp *op) {
+  auto const input1 = op->getInput();
+  auto const input2 = op->getStart();
+  auto const input3 = op->getSize();
+
+  auto input1ConstShape = dyn_cast<tosa::ConstShapeOp>(input1.getDefiningOp());
+
+  if (!input1ConstShape)
+    return {};
+
+  auto const input1Attr = cast<DenseElementsAttr>(input1ConstShape.getValues());
+  if (!input1Attr)
+    return {};
+
+  auto const input1Vals = input1Attr.getValues<APInt>();
+  auto const totalInput1 = input1Vals.size();
+
+  auto const start = getSingleI64From1ElementTensor(input2);
+  auto const size = getSingleI64From1ElementTensor(input3);
+
+  if (failed(start) || failed(size))
+    return {};
+
+  auto const startV = static_cast<int32_t>(start.value());
+  auto const sizeV = static_cast<int32_t>(size.value());
+
+  if ((sizeV <= 0) || (startV < 0) ||
+      (static_cast<size_t>(startV + sizeV) > totalInput1))
+    return {};
+
+  SmallVector<APInt> sliceOfInput;
+  sliceOfInput.reserve(totalInput1);
+
+  for (auto i = startV; i < (startV + sizeV); i++) {
+    sliceOfInput.push_back(input1Vals[i]);
+  }
----------------
lhutton1 wrote:

nit: could be simplified to?
```
SmallVector<APInt> sliceOfInput(
    input1Vals.begin() + startV,
    input1Vals.begin() + startV + sizeV
);
```

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