bug: (csl-lowering) Make multi-apply lowering work

tests/filecheck/transforms/convert-stencil-to-csl-stencil.mlir

@@ -238,8 +238,8 @@ builtin.module {
 // CHECK-NEXT:   %0 = tensor.empty() : tensor<1x64xf32>
 // CHECK-NEXT:   csl_stencil.apply(%arg1 : !stencil.field<[-1,1]x[-1,1]xtensor<64xf32>>, %0 : tensor<1x64xf32>) -> () <{"swaps" = [#csl_stencil.exchange<to [-1, 0]>], "topo" = #dmp.topo<64x64>, "num_chunks" = 2 : i64, "operandSegmentSizes" = array<i32: 1, 1, 0, 0, 0>}> ({
 // CHECK-NEXT:   ^0(%1 : tensor<1x32xf32>, %2 : index, %3 : tensor<1x64xf32>):
-// CHECK-NEXT:     %4 = csl_stencil.access %3[-1, 0] : tensor<1x64xf32>
-// CHECK-NEXT:     %5 = "tensor.insert_slice"(%4, %3, %2) <{"static_offsets" = array<i64: 0, -9223372036854775808>, "static_sizes" = array<i64: 32>, "static_strides" = array<i64: 1>, "operandSegmentSizes" = array<i32: 1, 1, 1, 0, 0>}> : (tensor<32xf32>, tensor<1x64xf32>, index) -> tensor<1x64xf32>
+// CHECK-NEXT:     %4 = csl_stencil.access %1[-1, 0] : tensor<1x32xf32>
+// CHECK-NEXT:     %5 = "tensor.insert_slice"(%4, %3, %2) <{"static_offsets" = array<i64: 0, -9223372036854775808>, "static_sizes" = array<i64: 1, 32>, "static_strides" = array<i64: 1, 1>, "operandSegmentSizes" = array<i32: 1, 1, 1, 0, 0>}> : (tensor<32xf32>, tensor<1x64xf32>, index) -> tensor<1x64xf32>
 // CHECK-NEXT:     csl_stencil.yield %5 : tensor<1x64xf32>
 // CHECK-NEXT:   }, {
 // CHECK-NEXT:   ^1(%6 : !stencil.field<[-1,1]x[-1,1]xtensor<64xf32>>, %7 : tensor<1x64xf32>):

xdsl/transforms/convert_stencil_to_csl_stencil.py

@@ -608,17 +608,17 @@ def match_and_rewrite(
         block2 = Block(arg_types=[op.input_stencil.type, op.result.type])
         block2.add_op(csl_stencil.YieldOp(block2.args[1]))
 
-        with ImplicitBuilder(block) as (_, offset, acc):
+        with ImplicitBuilder(block) as (buf, offset, acc):
             dest = acc
             for i, acc_offset in enumerate(offsets):
                 ac_op = csl_stencil.AccessOp(
-                    dest, stencil.IndexAttr.get(*acc_offset), chunk_t
+                    buf, stencil.IndexAttr.get(*acc_offset), chunk_t
                 )
                 assert isa(ac_op.result.type, AnyTensorType)
                 dest = tensor.InsertSliceOp.get(
                     source=ac_op.result,
                     dest=dest,
-                    static_sizes=ac_op.result.type.get_shape(),
+                    static_sizes=[1, *ac_op.result.type.get_shape()],
                     static_offsets=[i, memref.SubviewOp.DYNAMIC_INDEX],
                     offsets=[offset],
                 ).result

xdsl/transforms/csl_stencil_bufferize.py

@@ -94,6 +94,9 @@ def match_and_rewrite(self, op: csl_stencil.ApplyOp, rewriter: PatternRewriter,
         # convert args
         buf_args: list[SSAValue] = []
         to_memrefs: list[Operation] = [buf_iter_arg := to_memref_op(op.accumulator)]
+        op.accumulator.replace_by_if(
+            buf_iter_arg.memref, lambda use: use.operation != buf_iter_arg
+        )
         for arg in [*op.args_rchunk, *op.args_dexchng]:
             if isa(arg.type, TensorType[Attribute]):
                 to_memrefs.append(new_arg := to_memref_op(arg))
@@ -385,6 +388,9 @@ def match_and_rewrite(self, op: arith.ConstantOp, rewriter: PatternRewriter, /):
 class InjectApplyOutsIntoLinalgOuts(RewritePattern):
     @op_type_rewrite_pattern
     def match_and_rewrite(self, op: csl_stencil.ApplyOp, rewriter: PatternRewriter, /):
+        if not op.dest:
+            return
+
         yld = op.done_exchange.block.last_op
         assert isinstance(yld, csl_stencil.YieldOp)
         new_dest: list[SSAValue] = []

xdsl/transforms/experimental/stencil_tensorize_z_dimension.py

@@ -84,9 +84,13 @@ def get_required_result_type(op: Operation) -> TensorType[Attribute] | None:
                     tuple[int, ...],
                 )
             ):
+                assert is_tensor(use.operation.source.type)
+                # inserting a 1d tensor into a 2d array should not require the input tensor to also be 2d
+                # instead, drop the first `dimdiff` dimensions
+                dimdiff = len(static_sizes) - len(use.operation.source.type.shape)
                 return TensorType(
                     use.operation.result.type.get_element_type(),
-                    static_sizes,
+                    static_sizes[dimdiff:],
                 )
             for ret in use.operation.results:
                 if isa(r_type := ret.type, TensorType[Attribute]):