diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py index 14c7380e432f0..d9ab504f0de54 100644 --- a/mlir/python/mlir/dialects/transform/structured.py +++ b/mlir/python/mlir/dialects/transform/structured.py @@ -713,6 +713,8 @@ def __init__( disable_transfer_permutation_map_lowering_patterns: bool = False, vectorize_nd_extract: bool = False, vectorize_padding: bool = False, + flatten_1d_depthwise_conv: bool = False, + fold_type_extensions_into_contract: bool = False, loc=None, ip=None, ): @@ -722,8 +724,10 @@ def __init__( target, disable_multi_reduction_to_contract_patterns=disable_multi_reduction_to_contract_patterns, disable_transfer_permutation_map_lowering_patterns=disable_transfer_permutation_map_lowering_patterns, + flatten_1d_depthwise_conv=flatten_1d_depthwise_conv, vectorize_nd_extract=vectorize_nd_extract, vectorize_padding=vectorize_padding, + fold_type_extensions_into_contract=fold_type_extensions_into_contract, loc=loc, ip=ip, ) diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py index d6b70dc9d1978..e58b7646316fc 100644 --- a/mlir/test/python/dialects/transform_structured_ext.py +++ b/mlir/test/python/dialects/transform_structured_ext.py @@ -627,12 +627,16 @@ def testVectorizeChildrenAndApplyPatternsAllAttrs(target): disable_transfer_permutation_map_lowering_patterns=True, vectorize_nd_extract=True, vectorize_padding=True, + flatten_1d_depthwise_conv=True, + fold_type_extensions_into_contract=True, ) # CHECK-LABEL: TEST: testVectorizeChildrenAndApplyPatternsAllAttrs # CHECK: transform.sequence # CHECK: = transform.structured.vectorize # CHECK-SAME: disable_multi_reduction_to_contract_patterns # CHECK-SAME: disable_transfer_permutation_map_lowering_patterns + # CHECK-SAME: flatten_1d_depthwise_conv + # CHECK-SAME: fold_type_extensions_into_contract # CHECK-SAME: vectorize_nd_extract # CHECK-SAME: vectorize_padding @@ -646,12 +650,16 @@ def testVectorizeChildrenAndApplyPatternsNoAttrs(target): disable_transfer_permutation_map_lowering_patterns=False, vectorize_nd_extract=False, vectorize_padding=False, + flatten_1d_depthwise_conv=False, + fold_type_extensions_into_contract=False, ) # CHECK-LABEL: TEST: testVectorizeChildrenAndApplyPatternsNoAttrs # CHECK: transform.sequence # CHECK: = transform.structured.vectorize # CHECK-NOT: disable_multi_reduction_to_contract_patterns # CHECK-NOT: disable_transfer_permutation_map_lowering_patterns + # CHECK-NOT: flatten_1d_depthwise_conv + # CHECK-NOT: fold_type_extensions_into_contract # CHECK-NOT: vectorize_nd_extract # CHECK-NOT: vectorize_padding