Cose belle

src/PIM/Common/IR/TensorSliceUtils.cpp

@@ -68,4 +68,39 @@ Value insertStaticSlice(
     .getResult();
 }
 
+FailureOr<Value> addLeadingUnitTensorDimension(OpBuilder& builder, Location loc, Value value) {
+  auto type = dyn_cast<RankedTensorType>(value.getType());
+  if (!type || !type.hasStaticShape())
+    return failure();
+  SmallVector<int64_t> shape {1};
+  llvm::append_range(shape, type.getShape());
+  auto resultType = RankedTensorType::get(shape, type.getElementType(), type.getEncoding());
+  SmallVector<ReassociationIndices> reassociation;
+  if (type.getRank() != 0) {
+    reassociation.push_back({0, 1});
+    for (int64_t dim = 1; dim < type.getRank(); ++dim)
+      reassociation.push_back({dim + 1});
+  }
+  return tensor::ExpandShapeOp::create(builder, loc, resultType, value, reassociation).getResult();
+}
+
+FailureOr<Value> removeLeadingUnitTensorDimension(
+  OpBuilder& builder, Location loc, Value value, RankedTensorType resultType) {
+  if (value.getType() == resultType)
+    return value;
+  auto type = dyn_cast<RankedTensorType>(value.getType());
+  if (!type || !resultType || !type.hasStaticShape() || !resultType.hasStaticShape()
+      || type.getRank() != resultType.getRank() + 1 || type.getDimSize(0) != 1
+      || type.getElementType() != resultType.getElementType()
+      || !llvm::equal(type.getShape().drop_front(), resultType.getShape()))
+    return failure();
+  SmallVector<ReassociationIndices> reassociation;
+  if (resultType.getRank() != 0) {
+    reassociation.push_back({0, 1});
+    for (int64_t dim = 1; dim < resultType.getRank(); ++dim)
+      reassociation.push_back({dim + 1});
+  }
+  return tensor::CollapseShapeOp::create(builder, loc, resultType, value, reassociation).getResult();
+}
+
 } // namespace onnx_mlir