diff --git a/xdsl/transforms/csl_stencil_bufferize.py b/xdsl/transforms/csl_stencil_bufferize.py index bd4f46d9fe..add9764e36 100644 --- a/xdsl/transforms/csl_stencil_bufferize.py +++ b/xdsl/transforms/csl_stencil_bufferize.py @@ -289,10 +289,18 @@ def match_and_rewrite(self, op: csl_stencil.AccessOp, rewriter: PatternRewriter, rewriter.replace_matched_op(to_tensor_op(op.op)) return + # accesses to buffers passed in additional args can read directly from memref underlying `to_tensor` + source = ( + op.op.op.memref + if isinstance(op.op, OpResult) + and isinstance(op.op.op, bufferization.ToTensorOp) + else op.op + ) + rewriter.replace_matched_op( [ access := csl_stencil.AccessOp( - op.op, + source, op.offset, r_type, op.offset_mapping,