Skip to content

Commit

Permalink
[TIR] Fix printing enum in TransformLayout::AsPython (apache#11211)
Browse files Browse the repository at this point in the history
After this PR, `as_python` can handle `transform_layout` correctly. The result will be like

```python
sch.transform_layout(..., buffer_index_type="read", ...)
```

Previously `buffer_index_type` was printed unquoted.
  • Loading branch information
vinx13 authored May 5, 2022
1 parent aa3bcd9 commit 5007033
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions src/tir/schedule/primitive/layout_transformation.cc
Original file line number Diff line number Diff line change
Expand Up @@ -207,8 +207,10 @@ struct TransformLayoutTraits : public UnpackedInstTraits<TransformLayoutTraits>
PythonAPICall py("transform_layout");
py.Input("block", block_rv);
py.Input("buffer_index", buffer_index);
py.Input("buffer_index_type",
BufferIndexType2Str(static_cast<BufferIndexType>(buffer_index_type->value)));
py.Input("buffer_index_type", '"' +
std::string(BufferIndexType2Str(
static_cast<BufferIndexType>(buffer_index_type->value))) +
'"');
py.Input("index_map", index_map->ToPythonString());
return py.Str();
}
Expand Down

0 comments on commit 5007033

Please sign in to comment.