Adddress some comments

Author: Mr-Anyone

mlir/lib/Conversion/LinalgToRock/LinalgToRock.cpp

@@ -169,6 +169,27 @@ LogicalResult ExpandStrideConverter::matchAndRewrite(
     return failure();
   }
 
+  // Require statically known slice sizes that exactly match the
+  // source tensor shape.
+  auto srcType = dyn_cast<RankedTensorType>(op.getSource().getType());
+  if (!srcType){
+    return failure();
+  }
+
+  // into rock.expand_strides, but only in the exact expand-strides shape:
+  // - dest is a tensor.empty with a single use
+  // - all offsets are zero
+  // - all strides are one
+  // - all slice sizes are static and match the source tensor shape
+  bool isExpandStride = llvm::all_of(op.getStaticOffsets(), [](int64_t offset) { return offset == 0; }) &&
+    llvm::all_of(op.getStaticStrides(), [](int64_t stride) { return stride == 1; }) &&
+    llvm::none_of(op.getStaticSizes(),
+        [](int64_t s) { return s == ShapedType::kDynamic; }) && op.getStaticSizes() == srcType.getShape();
+
+  if(!isExpandStride){
+    return failure();
+  }
+
   Location loc = op.getLoc();
   auto alloc = bufferization::AllocTensorOp::create(
       rewriter, loc, tensorEmpty.getResult().getType(), {});

mlir/lib/Conversion/LinalgToRock/LinalgToRockPass.cpp

@@ -43,7 +43,26 @@ static void populateLinalgToRockDialectConversion(ConversionTarget &target) {
   target.addDynamicallyLegalOp<tensor::InsertSliceOp>(
       [](tensor::InsertSliceOp op) -> std::optional<bool> {
         auto emptyOp = op.getDest().getDefiningOp<tensor::EmptyOp>();
-        return !(emptyOp && emptyOp->hasOneUse());
+        if (!emptyOp){
+          return true;
+        }
+
+        // Require statically known slice sizes that exactly match the
+        // source tensor shape.
+        auto srcType = dyn_cast<RankedTensorType>(op.getSource().getType());
+        if (!srcType)
+          return true;
+
+        // into rock.expand_strides, but only in the exact expand-strides shape:
+        // - dest is a tensor.empty with a single use
+        // - all offsets are zero
+        // - all strides are one
+        // - all slice sizes are static and match the source tensor shape
+        bool isExpandStride = llvm::all_of(op.getStaticOffsets(), [](int64_t offset) { return offset == 0; }) &&
+          llvm::all_of(op.getStaticStrides(), [](int64_t stride) { return stride == 1; }) &&
+          llvm::none_of(op.getStaticSizes(),
+                  [](int64_t s) { return s == ShapedType::kDynamic; }) && op.getStaticSizes() == srcType.getShape();
+        return !isExpandStride;
       });
 
   // We only allow Linalg operations that are elementwise. Fusion is supported

mlir/test/Conversion/MIGraphXToLinalg/mixr-to-linalg-ops.mlir

@@ -144,14 +144,12 @@ func.func @clip_i32(%arg0: !migraphx.shaped<64x64xi32, 64x1>, %arg1: !migraphx.s
 // Test taken from migraphx-to-tosa
 // Tests for non-standard shapes.
 
-
-
 // CHECK-LABEL: func.func @transposed(
 func.func @transposed(%arg0: !migraphx.shaped<4x3xf32, 1x4>) -> !migraphx.shaped<4x3xf32, 1x4> {
   %op = migraphx.floor %arg0 : <4x3xf32, 1x4> -> <4x3xf32, 1x4>
   // CHECK: %[[FLOOR:.*]] = linalg.floor ins{{.*}} -> tensor<4x3xf32>
   // CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<3x4xf32>
-  // CHECK: linalg.transpose ins(%[[FLOOR]] : tensor{{.*}}) outs(%3 : tensor{{.*}}) permutation = [1, 0] 
+  // CHECK: linalg.transpose ins(%[[FLOOR]] : tensor{{.*}}) outs(%[[EMPTY]] : tensor{{.*}}) permutation = [1, 0] 
   func.return %op : !migraphx.shaped<4x3xf32, 1x4>
 }
 
@@ -163,7 +161,7 @@ func.func @broadcast(%arg0: !migraphx.shaped<4x3xf32, 1x0>, %arg1: !migraphx.sha
   // CHECK-DAG:     %[[collapsed:.*]] = tensor.collapse_shape %[[expanded_0]]
   // CHECK-DAG:     %[[expanded_1:.*]] = tensor.expand_shape %[[collapsed]]
   // CHECK-DAG:     %[[zero:.*]] = tensor.empty() : tensor<4x3xf32>
-  // CHECK-DAG:     %[[broadcasted:.*]] = linalg.broadcast ins(%[[expanded_1]] : tensor<4xf32>) outs(%[[zero]] : tensor<4x3xf32>) dimensions = [1] 
+  // CHECK-DAG:     %[[broadcasted:.*]] = linalg.broadcast ins(%[[expanded_1]] : tensor<4xf32>) outs(%[[zero]] : tensor<4x3xf32>) dimensions = [1]
   %op = migraphx.sub %arg0, %arg1 : <4x3xf32, 1x0>, <4x3xf32, 3x1> -> <4x3xf32, 3x1>
   // CHECK-DAG:     %[[one:.*]] = tensor.empty() : tensor<4x3xf32>
   // CHECK-DAG:     %[[two:.*]] = linalg.sub ins(%[[broadcasted]], %[[expanded]] : tensor<4x3xf32>, tensor<4x3xf32>) outs(%[[one]] : tensor<4x3xf32>) -> tensor<4x3xf32>