Hzfengsy commented on code in PR #12720:
URL: https://github.com/apache/tvm/pull/12720#discussion_r964422619
##########
include/tvm/tir/schedule/schedule.h:
##########
@@ -601,9 +601,11 @@ class ScheduleNode : public runtime::Object {
* \param buffer_index The index of the buffer in block's read or write
region.
* \param buffer_index_type The type of the buffer index, kRead or kWrite.
* \param index_map The transformation to apply.
+ * \param pad_value The value to write into padding introduced by the
transformation.
*/
virtual void TransformLayout(const BlockRV& block_rv, int buffer_index,
- BufferIndexType buffer_index_type, const
IndexMap& index_map) = 0;
+ BufferIndexType buffer_index_type, const
IndexMap& index_map,
+ const Optional<PrimExpr>& pad_value) = 0;
Review Comment:
How about adding a default value?
```suggestion
const Optional<PrimExpr>& pad_value =
NullOpt) = 0;
```
##########
tests/python/unittest/test_tir_schedule_transform_layout.py:
##########
@@ -329,5 +329,302 @@ def
test_transform_block_layout_fail_mixed_iter_type(use_block_name):
)
+class BasePaddingCompare(tvm.testing.CompareBeforeAfter):
+ pad_value = tvm.testing.parameter(None)
+
+ transformed_buffer = tvm.testing.parameter("A")
+
+ @pytest.fixture
+ def transform(self, pad_value, transformed_buffer):
+ def transform(mod):
+ sch = tir.Schedule(mod)
+ sch.transform_layout(
+ "block", transformed_buffer, lambda i: [i // 4, i % 4],
pad_value=pad_value
+ )
+ # sch.transform_block_layout("block", lambda i: [i // 4, i % 4])
+ return sch.mod
+
+ return transform
+
+
+class TestNoPadding(BasePaddingCompare):
+ """Transformations without padding do not depend on pad_value."""
+
+ pad_value = tvm.testing.parameter(None, 42)
+
+ def before():
Review Comment:
It's great if we can support padding opaque blocks. However, non-opaque
blocks are the most common cases.
Could you please change the test cases into non-opaque blocks?
e.g.
```python
with T.block():
vi = T.axis.remap("S", [i])
A[vi] = 0
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
