Skip to content

Commit

Permalink
Third attempt to merge #191 (#307)
Browse files Browse the repository at this point in the history
* Changes to `_InlineWindow`

* Changes to `_CallSwap`

* `Check_Aliasing` in `_BindConfig` and `_BindExpr`

* Changes to `_DoExpandDim`

* Changes to `_DoRearrangeDim`

* Changes to `_DoRemoveLoop`

* Changes to `_DoFissionAfterSimple`

* Changes to `_DoFuseIf`

* Added `fuse_loop` tests

* Changes to `_DoAddLoop`

* Changes to `_DoExtractMethod`

* Changes to `_AssertIf`

* Changes to `_DoDataReuse`

* Changes to `_DoStageWindow`

* Changes to `_DoBoundAlloc`

* Changes to `DoReplace`

* Update commented line in `_DoAddUnsafeGuard`

* fixed AMX test error

* Update src/exo/LoopIR_scheduling.py

---------

Co-authored-by: Kevin Qian <keqian@mit.edu>
Co-authored-by: Yuka Ikarashi <yukatkh@gmail.com>
  • Loading branch information
3 people committed Jan 31, 2023
1 parent 86f2a15 commit 6cc4596
Show file tree
Hide file tree
Showing 14 changed files with 676 additions and 241 deletions.
14 changes: 7 additions & 7 deletions src/exo/API_scheduling.py
Original file line number Diff line number Diff line change
Expand Up @@ -1142,30 +1142,30 @@ def expand_dim(proc, buf_cursor, alloc_dim, indexing_expr, unsafe_disable_checks


@sched_op([AllocCursorA, ListA(IntA)])
def rearrange_dim(proc, buf_cursor, dimensions):
def rearrange_dim(proc, buf_cursor, permute_vector):
"""
Rearranges the dimensions of the indicated buffer allocation according
to the supplied permutation (`dimensions`).
to the supplied permutation (`permute_vector`).
args:
buf_cursor - cursor pointing to an Alloc statement
for an N-dimensional array
dimensions - a permutation of the integers (0,1,...,N-1)
permute_vector - a permutation of the integers (0,1,...,N-1)
rewrite:
(with dimensions = [2,0,1])
(with permute_vector = [2,0,1])
`x : T[N,M,K]` -> `x : T[K,N,M]`
"""
proc_c = ic.Cursor.root(proc)
stmt = buf_cursor._impl
# extra sanity check
N = len(stmt._node().type.hi)
if set(range(0, N)) != set(dimensions):
if set(range(0, N)) != set(permute_vector):
raise ValueError(
f"dimensions argument ({dimensions}) "
f"permute_vector argument ({permute_vector}) "
f"was not a permutation of {set(range(0, N))}"
)
return Schedules.DoRearrangeDim(proc_c, stmt, dimensions).result()
return Schedules.DoRearrangeDim(proc_c, stmt, permute_vector).result()


@sched_op([AllocCursorA, ListA(OptionalA(NewExprA("buf_cursor"))), BoolA])
Expand Down
Loading

0 comments on commit 6cc4596

Please sign in to comment.