Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

dialects: stencil: Allow loading and storing from a same field. #3071

Merged
merged 7 commits into from
Aug 21, 2024
Merged

dialects: stencil: Allow loading and storing from a same field. #3071

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
62d5360
Move "cannot load and store" constraint from the stencil dialect to t…
PapyChacal Aug 20, 2024
5ff4211
Tests plz
PapyChacal Aug 20, 2024
21740b0
Merge branch 'main' into emilien/stencil-load-store
PapyChacal Aug 20, 2024
cd91cbf
Merge branch 'main' into emilien/stencil-load-store
PapyChacal Aug 20, 2024
43e34a8
Merge branch 'main' into emilien/stencil-load-store
PapyChacal Aug 20, 2024
935675e
Merge branch 'main' into emilien/stencil-load-store
PapyChacal Aug 20, 2024
4d35c83
Merge branch 'main' into emilien/stencil-load-store
PapyChacal Aug 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 0 additions & 21 deletions tests/dialects/test_stencil.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -529,27 +529,6 @@ def test_stencil_store():
assert store.bounds is bounds


def test_stencil_store_load_overlap():
temp_type = TempType([(0, 5), (0, 5)], f32)
temp_type_ssa_val = TestSSAValue(temp_type)

field_type = FieldType([(0, 2), (0, 2)], f32)
field_type_ssa_val = TestSSAValue(field_type)

lb = IndexAttr.get(1, 1)
ub = IndexAttr.get(64, 64)
bounds = StencilBoundsAttr.new((lb, ub))

load = LoadOp.get(field_type_ssa_val, lb, ub)
store = StoreOp.get(temp_type_ssa_val, field_type_ssa_val, bounds)

with pytest.raises(VerifyException, match="Cannot Load and Store the same field!"):
load.verify()

with pytest.raises(VerifyException, match="Cannot Load and Store the same field!"):
store.verify()


PapyChacal marked this conversation as resolved.
Show resolved Hide resolved
def test_stencil_index():
dim = IntAttr(10)
offset = IndexAttr.get(1)
Expand Down
53 changes: 52 additions & 1 deletion tests/filecheck/transforms/convert-stencil-to-ll-mlir.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// RUN: xdsl-opt %s -p convert-stencil-to-ll-mlir | filecheck %s
// RUN: xdsl-opt %s -p convert-stencil-to-ll-mlir --split-input-file --verify-diagnostics | filecheck %s

builtin.module {
// CHECK: builtin.module {
Expand Down Expand Up @@ -844,3 +844,54 @@ func.func @buffered_combine(%115 : !stencil.field<?x?xf64>) {

}
// CHECK-NEXT: }

// -----

func.func @stencil_double_store(%93 : !stencil.field<[0,64]x[0,64]x[0,64]xindex>) {
%94 = stencil.apply() -> (!stencil.temp<[0,64]x[0,64]x[0,64]xindex>) {
%x_1 = stencil.index 0 <[1, 1, 1]>
%y_1 = stencil.index 1 <[-1, -1, -1]>
%z_1 = stencil.index 2 <[0, 0, 0]>
%xy_1 = arith.addi %x_1, %y_1 : index
%xyz_1 = arith.addi %xy_1, %z_1 : index
stencil.return %xyz_1 : index
}
stencil.store %94 to %93 (<[0, 0, 0], [64, 64, 64]>) : !stencil.temp<[0,64]x[0,64]x[0,64]xindex> to !stencil.field<[0,64]x[0,64]x[0,64]xindex>
%95 = stencil.apply() -> (!stencil.temp<[0,64]x[0,64]x[0,64]xindex>) {
%x_1 = stencil.index 0 <[1, 1, 1]>
%y_1 = stencil.index 1 <[-1, -1, -1]>
%z_1 = stencil.index 2 <[0, 0, 0]>
%xy_1 = arith.addi %x_1, %y_1 : index
%xyz_1 = arith.addi %xy_1, %z_1 : index
stencil.return %xyz_1 : index
}
stencil.store %95 to %93 (<[0, 0, 0], [64, 64, 64]>) : !stencil.temp<[0,64]x[0,64]x[0,64]xindex> to !stencil.field<[0,64]x[0,64]x[0,64]xindex>
func.return
}

// CHECK: "stencil.store"(%2, %0) {"bounds" = #stencil.bounds<[0, 0, 0], [64, 64, 64]>} : (!stencil.temp<[0,64]x[0,64]x[0,64]xindex>, !stencil.field<[0,64]x[0,64]x[0,64]xindex>) -> ()
// CHECK-NEXT: ^^^^^^^^^^^^^^^-------------------------------------------------------------------------------------------------------------------------
// CHECK-NEXT: | Error while applying pattern: Cannot lower directly if storing to the same field multiple times! Try running `stencil-bufferize` before.
// CHECK-NEXT: ----------------------------------------------------------------------------------------------------------------------------------------
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so, we still keep this, as expected

Copy link
Collaborator Author

@PapyChacal PapyChacal Aug 21, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Kinda, as a transition; This is now only an error to try and use the old lowering on such a case!
So you can now express it without hacks, and in most cases lower it simply by first running stencil-bufferize

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(Where before, it was considered an error to just write it!)


// -----

func.func @stencil_load_store(%93 : !stencil.field<[0,64]x[0,64]x[0,64]xindex>) {
%94 = stencil.apply() -> (!stencil.temp<[0,64]x[0,64]x[0,64]xindex>) {
%x_1 = stencil.index 0 <[1, 1, 1]>
%y_1 = stencil.index 1 <[-1, -1, -1]>
%z_1 = stencil.index 2 <[0, 0, 0]>
%xy_1 = arith.addi %x_1, %y_1 : index
%xyz_1 = arith.addi %xy_1, %z_1 : index
stencil.return %xyz_1 : index
}
stencil.store %94 to %93 (<[0, 0, 0], [64, 64, 64]>) : !stencil.temp<[0,64]x[0,64]x[0,64]xindex> to !stencil.field<[0,64]x[0,64]x[0,64]xindex>
%95 = stencil.load %93 : !stencil.field<[0,64]x[0,64]x[0,64]xindex> -> !stencil.temp<[0,64]x[0,64]x[0,64]xindex>
"test.op"(%95) : (!stencil.temp<[0,64]x[0,64]x[0,64]xindex>) -> ()
func.return
}

// CHECK: %2 = "stencil.load"(%0) : (!stencil.field<[0,64]x[0,64]x[0,64]xindex>) -> !stencil.temp<[0,64]x[0,64]x[0,64]xindex>
// CHECK-NEXT: ^^^^^^^^^^^^^^^^^^^-----------------------------------------------------------------------------------------------------------------
// CHECK-NEXT: | Error while applying pattern: Cannot lower directly if loading and storing the same field! Try running `stencil-bufferize` before.
// CHECK-NEXT: ------------------------------------------------------------------------------------------------------------------------------------
10 changes: 0 additions & 10 deletions xdsl/dialects/stencil.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1277,9 +1277,6 @@ def get(
)

def verify_(self) -> None:
for use in self.field.uses:
if isa(use.operation, StoreOp):
raise VerifyException("Cannot Load and Store the same field!")
field = self.field.type
temp = self.res.type
assert isa(field, FieldType[Attribute])
Expand Down Expand Up @@ -1467,13 +1464,6 @@ def get(
):
return StoreOp.build(operands=[temp, field], attributes={"bounds": bounds})

def verify_(self) -> None:
for use in self.field.uses:
if isa(use.operation, LoadOp):
raise VerifyException("Cannot Load and Store the same field!")
if isa(use.operation, LoadOp) and use.operation is not self:
raise VerifyException("Can only store once to a field!")


@irdl_op_definition
class StoreResultOp(IRDLOperation):
Expand Down
14 changes: 14 additions & 0 deletions xdsl/transforms/experimental/convert_stencil_to_ll_mlir.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,11 @@ def assert_subset(field: FieldType[Attribute], temp: TempType[Attribute]):
class LoadOpToMemref(RewritePattern):
@op_type_rewrite_pattern
def match_and_rewrite(self, op: LoadOp, rewriter: PatternRewriter, /):
for use in op.field.uses:
if isa(use.operation, StoreOp):
raise VerifyException(
"Cannot lower directly if loading and storing the same field! Try running `stencil-bufferize` before."
)
field = op.field.type
assert isa(field, FieldType[Attribute])
assert isa(field.bounds, StencilBoundsAttr)
Expand Down Expand Up @@ -501,6 +506,15 @@ class StencilStoreToSubview(RewritePattern):

@op_type_rewrite_pattern
def match_and_rewrite(self, op: StoreOp, rewriter: PatternRewriter, /):
for use in op.field.uses:
if isa(use.operation, LoadOp):
raise VerifyException(
"Cannot lower directly if loading and storing the same field! Try running `stencil-bufferize` before."
)
if isa(use.operation, StoreOp) and use.operation is not op:
raise VerifyException(
"Cannot lower directly if storing to the same field multiple times! Try running `stencil-bufferize` before."
)
field = op.field
assert isa(field.type, FieldType[Attribute])
assert isa(field.type.bounds, StencilBoundsAttr)
Expand Down
Loading