Resnet is fast

2026-07-20 11:34:59 +02:00
parent 5f42da36ae
commit 6bad9a8008
16 changed files with 883 additions and 138 deletions
@@ -80,8 +80,17 @@ Value extractMixedSliceOrIdentity(RewriterBase &rewriter,

 Value insertMixedSlice(OpBuilder &builder, Location loc, Value source,
                       Value dest, const MixedSliceGeometry &geometry) {
+  SmallVector<OpFoldResult> sizes(geometry.sizes);
+  auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+  auto destType = dyn_cast<RankedTensorType>(dest.getType());
+  if (sourceType && destType && sourceType.hasStaticShape()
+      && sourceType.getRank() == destType.getRank()) {
+    sizes.clear();
+    for (int64_t dimension : sourceType.getShape())
+      sizes.push_back(builder.getIndexAttr(dimension));
+  }
  return tensor::InsertSliceOp::create(builder, loc, source, dest,
-                                       geometry.offsets, geometry.sizes,
+                                       geometry.offsets, sizes,
                                       geometry.strides);
 }