Drop some patterns from IREE apply_patterns op (#14053) Drop these patterns from IREE the `apply_patterns` op and use the upstreamed versions: `fold_memref_aliases`, `fold_tensor_subsets`, `rank_reducing_linalg`, `rank_reducing_linalg_via_reshapes`, `rank_reducing_vector`
diff --git a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp index c3ac3b2..9d0468f 100644 --- a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp
@@ -184,6 +184,11 @@ patterns.insert<FoldFillIntoPad>(patterns.getContext()); } +void transform_dialect::ApplyFoldReshapeIntoTensorHalInterfacePatternsOp:: + populatePatterns(RewritePatternSet &patterns) { + populateReshapeToInterfaceTensorPatterns(patterns); +} + //===---------------------------------------------------------------------===// // ApplyPatternsOp //===---------------------------------------------------------------------===// @@ -211,11 +216,7 @@ ADD_PATTERN(expandMemrefStridedMetadata, getExpandMemrefStridedMetadataAttrName) ADD_PATTERN(extractAddressComputations, getExtractAddressComputationsAttrName) - ADD_PATTERN(foldMemrefAliases, getFoldMemrefAliasesAttrName) ADD_PATTERN(foldReassociativeReshapes, getFoldReassociativeReshapesAttrName) - ADD_PATTERN(foldTensorSubsets, getFoldTensorSubsetsAttrName) - ADD_PATTERN(foldVectorTransferTensorSlice, - getFoldVectorTransferTensorSliceAttrName) ADD_PATTERN(licm, getLicmAttrName) ADD_PATTERN(linalgElementwiseGreedyFusion, getLinalgElementwiseGreedyFusionAttrName) @@ -223,10 +224,6 @@ getLowerTransferOpPermutationsAttrName) ADD_PATTERN(lowerVectorMasks, getLowerVectorMasksAttrName) ADD_PATTERN(prepareVectorToMma, getPrepareVectorToMmaAttrName) - ADD_PATTERN(rankReducingLinalg, getRankReducingLinalgAttrName) - ADD_PATTERN(rankReducingLinalgViaReshapes, - getRankReducingLinalgViaReshapesAttrName) - ADD_PATTERN(rankReducingVector, getRankReducingVectorAttrName) ADD_PATTERN(swapPaddingElideConditional, getSwapPaddingElideConditionalAttrName) ADD_PATTERN(swappingPatterns, getSwappingPatternsAttrName) @@ -273,25 +270,10 @@ memref::populateExtractAddressComputationsPatterns(patterns); } -static void addFoldMemrefAliasPatterns(RewritePatternSet &patterns) { - memref::populateFoldMemRefAliasOpPatterns(patterns); -} - static void addReassociativeReshapePatterns(RewritePatternSet &patterns) { tensor::populateReassociativeReshapeFoldingPatterns(patterns); } -static void addFoldTensorSubsetsPatterns(RewritePatternSet &patterns) { - tensor::populateFoldTensorSubsetOpPatterns(patterns); - // TODO: upstream should move these to populateFoldTensorSubsetOpPatterns. - tensor::populateMergeConsecutiveInsertExtractSlicePatterns(patterns); -} - -static void addFoldVectorTransferTensorExtractPatterns( - RewritePatternSet &patterns) { - vector::populateVectorTransferTensorSliceTransforms(patterns); -} - static void addEraseUnnecessaryTensorOperandsPatterns( RewritePatternSet &patterns) { linalg::populateEraseUnnecessaryInputsPatterns(patterns); @@ -301,21 +283,6 @@ populatePrepareVectorToMMAPatterns(patterns, /*useNvGpu=*/true); } -static void addRankReducingLinalgPatterns(RewritePatternSet &patterns) { - populateReshapeToInterfaceTensorPatterns(patterns); - linalg::populateFoldUnitExtentDimsViaSlicesPatterns(patterns); -} - -static void addRankReducingLinalgViaReshapesPatterns( - RewritePatternSet &patterns) { - populateReshapeToInterfaceTensorPatterns(patterns); - linalg::populateFoldUnitExtentDimsViaReshapesPatterns(patterns); -} - -static void addRankReducingVectorPatterns(RewritePatternSet &patterns) { - vector::populateCastAwayVectorLeadingOneDimPatterns(patterns); -} - static void addSwappingPatterns(RewritePatternSet &patterns, bool swapPaddingElideCornerCase) { patterns.add<linalg::ExtractSliceOfPadTensorSwapPattern>( @@ -397,11 +364,7 @@ memref::populateExpandStridedMetadataPatterns(patterns); if (getExtractAddressComputations()) addExtractAddressComputationsPatterns(patterns); - if (getFoldMemrefAliases()) addFoldMemrefAliasPatterns(patterns); if (getFoldReassociativeReshapes()) addReassociativeReshapePatterns(patterns); - if (getFoldTensorSubsets()) addFoldTensorSubsetsPatterns(patterns); - if (getFoldVectorTransferTensorSlice()) - addFoldVectorTransferTensorExtractPatterns(patterns); if (getLinalgElementwiseGreedyFusion()) linalg::populateElementwiseOpsFusionPatterns(patterns, setFusedOpOperandLimit<3>); @@ -409,10 +372,6 @@ addLowerTransferOpPermutationsPatterns(patterns); if (getLowerVectorMasks()) addLowerVectorMasksPatterns(patterns); if (getPrepareVectorToMma()) addPrepareVectorToMmaPatterns(patterns); - if (getRankReducingLinalg()) addRankReducingLinalgPatterns(patterns); - if (getRankReducingLinalgViaReshapes()) - addRankReducingLinalgViaReshapesPatterns(patterns); - if (getRankReducingVector()) addRankReducingVectorPatterns(patterns); if (getSwappingPatterns()) addSwappingPatterns(patterns, getSwapPaddingElideConditional()); if (getUnrollVectorsGpuMmaSync())
diff --git a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h index 98e3a91..da6cada 100644 --- a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h +++ b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h
@@ -42,18 +42,12 @@ bool eraseUnnecessaryTensorOperands = false; bool expandMemrefStridedMetadata = false; bool extractAddressComputations = false; - bool foldMemrefAliases = false; bool foldReassociativeReshapes = false; - bool foldTensorSubsets = false; - bool foldVectorTransferTensorSlice = false; bool licm = false; bool linalgElementwiseGreedyFusion = false; bool lowerTransferOpPermutations = false; bool lowerVectorMasks = false; bool prepareVectorToMma = false; - bool rankReducingLinalg = false; - bool rankReducingLinalgViaReshapes = false; - bool rankReducingVector = false; bool swapPaddingElideConditional = false; bool swappingPatterns = false; bool unrollVectorsGpuMmaSync = false;
diff --git a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td index 33bd021..8103b01 100644 --- a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td +++ b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td
@@ -65,6 +65,18 @@ let assemblyFormat = "attr-dict"; } +def ApplyFoldReshapeIntoTensorHalInterfacePatternsOp : Op<Transform_Dialect, + "apply_patterns.iree.fold_reshape_into_tensor_hal_interface", + [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> { + let description = [{ + Populate patterns that fold tensor.expand_shape/tensor.collapse_shape into + the source hal.interface.binding.subspan op. + }]; + + let cppNamespace = "mlir::iree_compiler::IREE::transform_dialect"; + let assemblyFormat = "attr-dict"; +} + def ApplyPatternsOp : Op<Transform_Dialect, "iree.apply_patterns", [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>, TransformEachOpTrait, @@ -105,12 +117,8 @@ of their effect on the metadata (sizes, offset, strides). - extract_address_computations: adds patterns for anchoring subview accessing operations at [0, ... 0]. - - fold_memref_aliases: adds patterns for folding ops such as - memref.subview. - fold_reassociative_reshapes: adds patterns that fold insert_slice/ extract_slice ops with reassociative reshape ops. - - fold_tensor_subsets: adds patterns for folding tensor subset ops into - their producer and consumers. - licm: additionally apply loop-independent code motion and single iteration loop promotion. This is not a set of patterns per se but is still very convenient to apply it close to canonicalization and other greedy @@ -122,13 +130,7 @@ - lower_vector_masks: Lower vector.mask ops away. - prepare_vector_to_mma: pre-process vector.contract op to set it in a form that can be mapped to nvgpu.mma operations. - - rank_reducing_linalg: adds patterns that results in rank-reducing behavior on subset-based linalg operations using insert/extract slices. - - rank_reducing_linalg_via_reshapes: adds patterns that results in rank-reducing - behavior on subset-based linalg operations using expand/collapse shape ops. - - rank_reducing_vector: adds patterns that results in rank-reducing - behavior on subset-based vector operations. - adopts the upstream version. - swapping_patterns: adds patterns that swap operations for a better outcome. This is a catch all that can be refined further if/when needed. - swap_padding_elide_conditional: refines the tensor.pad + @@ -167,18 +169,12 @@ UnitAttr:$erase_unnecessary_tensor_operands, UnitAttr:$expand_memref_strided_metadata, UnitAttr:$extract_address_computations, - UnitAttr:$fold_memref_aliases, UnitAttr:$fold_reassociative_reshapes, - UnitAttr:$fold_tensor_subsets, - UnitAttr:$fold_vector_transfer_tensor_slice, UnitAttr:$licm, UnitAttr:$linalg_elementwise_greedy_fusion, UnitAttr:$lower_transfer_op_permutations, UnitAttr:$lower_vector_masks, UnitAttr:$prepare_vector_to_mma, - UnitAttr:$rank_reducing_linalg, - UnitAttr:$rank_reducing_linalg_via_reshapes, - UnitAttr:$rank_reducing_vector, UnitAttr:$swap_padding_elide_conditional, UnitAttr:$swapping_patterns, UnitAttr:$unroll_vectors_gpu_mma_sync,
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/reductions_codegen_spec.mlir b/compiler/src/iree/compiler/Codegen/Common/test/reductions_codegen_spec.mlir index 81937f5..6d5a0fa 100644 --- a/compiler/src/iree/compiler/Codegen/Common/test/reductions_codegen_spec.mlir +++ b/compiler/src/iree/compiler/Codegen/Common/test/reductions_codegen_spec.mlir
@@ -60,7 +60,11 @@ // Step 3. Rank-reduce. // =========================================================================== - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op // We don't perform any following transformation (vectorization, bufferizaton, // mapping) because this schedule is applied to Linalg-only code without the
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/attention.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/attention.mlir index bd6af02..22642a2 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/attention.mlir +++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/attention.mlir
@@ -53,7 +53,11 @@ // Vectorize function // ========================================== %func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %func_3 = transform.structured.vectorize %func : (!transform.any_op) -> !transform.any_op // Bufferization
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/set_transform_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/set_transform_strategy.mlir index cf2ec79..e69dc5d 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/set_transform_strategy.mlir +++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/set_transform_strategy.mlir
@@ -72,13 +72,17 @@ // CHECK: transform.iree.forall_to_workgroup %{{.*}} // CHECK: transform.iree.map_nested_forall_to_gpu_threads %{{.*}} workgroup_dims = [64, 2, 1] warp_dims = [2, 2, 1] // CHECK: transform.iree.hoist_static_alloc %{{.*}} -// CHECK: transform.iree.apply_patterns %{{.*}} {fold_memref_aliases} +// CHECK: apply_patterns to %{{.*}} { +// CHECK: transform.apply_patterns.memref.fold_memref_alias_ops +// CHECK: } : !transform.any_op // CHECK: transform.iree.apply_patterns %{{.*}} {extract_address_computations} // CHECK: transform.iree.apply_patterns %{{.*}} {unroll_vectors_gpu_wmma} // CHECK: transform.structured.hoist_redundant_vector_transfers %{{.*}} // CHECK: transform.iree.apply_buffer_optimizations %{{.*}} // CHECK: transform.iree.vector.vector_to_mma_conversion %{{.*}} {use_wmma} -// CHECK: transform.iree.apply_patterns %{{.*}} {fold_memref_aliases} +// CHECK: apply_patterns to %{{.*}} { +// CHECK: transform.apply_patterns.memref.fold_memref_alias_ops +// CHECK: } : !transform.any_op // CHECK: transform.memref.multibuffer %{{.*}} {factor = 3 : i64, skip_analysis} // CHECK: transform.apply_patterns.vector.transfer_to_scf max_transfer_rank = 1 full_unroll = true // CHECK: transform.iree.create_async_groups %{{.*}} {use_mma_sync = false} @@ -125,7 +129,9 @@ // The warp dimensions are controled by td-matmul-strategy-num-warps-XX. // WITH_OPTIONS: transform.iree.map_nested_forall_to_gpu_threads %{{.*}} workgroup_dims = [32, 4, 1] warp_dims = [1, 4, 1] // WITH_OPTIONS: transform.iree.hoist_static_alloc %{{.*}} -// WITH_OPTIONS: transform.iree.apply_patterns %{{.*}} {fold_memref_aliases} +// WITH_OPTIONS: apply_patterns to %{{.*}} { +// WITH_OPTIONS: transform.apply_patterns.memref.fold_memref_alias_ops +// WITH_OPTIONS: } : !transform.any_op // WITH_OPTIONS: transform.iree.apply_patterns %{{.*}} {extract_address_computations} // The unroll attribute should match td-matmul-use-mma-sync, for true: mma_sync, // for false:_wmma. @@ -134,7 +140,9 @@ // WITH_OPTIONS: transform.iree.apply_buffer_optimizations %{{.*}} // The attribute should match td-matmul-use-mma-sync. // WITH_OPTIONS: transform.iree.vector.vector_to_mma_conversion %{{.*}} {use_mma_sync} -// WITH_OPTIONS: transform.iree.apply_patterns %{{.*}} {fold_memref_aliases} +// WITH_OPTIONS: apply_patterns to %{{.*}} { +// WITH_OPTIONS: transform.apply_patterns.memref.fold_memref_alias_ops +// WITH_OPTIONS: } : !transform.any_op // The multibuffer pass is only run when we set use-async-copies. // The factor should match td-matmul-strategy-pipeline-depth: 5. // WITH_OPTIONS: transform.memref.multibuffer %{{.*}} {factor = 5 : i64, skip_analysis}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/set_transform_strategy_pad.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/set_transform_strategy_pad.mlir index d739675..bb29693 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/set_transform_strategy_pad.mlir +++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/set_transform_strategy_pad.mlir
@@ -59,7 +59,11 @@ // CHECK: transform.structured.masked_vectorize {{.*}} vector_sizes [4, 4] : !transform.any_op // CHECK: {{.*}} = transform.structured.match ops{["func.func"]} in {{.*}} : (!transform.any_op) -> !transform.any_op // CHECK: transform.apply_patterns.vector.lower_masked_transfers -// CHECK: transform.iree.apply_patterns {{.*}} {rank_reducing_linalg, rank_reducing_vector} : (!transform.any_op) -> () +// CHECK: apply_patterns to %{{.*}} { +// CHECK-DAG: transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface +// CHECK-DAG: transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices +// CHECK-DAG: transform.apply_patterns.vector.cast_away_vector_leading_one_dim +// CHECK: } : !transform.any_op // CHECK: {{.*}} = transform.structured.vectorize {{.*}} : (!transform.any_op) -> !transform.any_op // CHECK: transform.iree.apply_patterns {{.*}} {canonicalization, cse, licm} : (!transform.any_op) -> () // CHECK: transform.iree.eliminate_empty_tensors {{.*}} : (!transform.any_op) -> ()
diff --git a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.cpp b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.cpp index 204c560..4e4fff9 100644 --- a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.cpp +++ b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.cpp
@@ -296,12 +296,13 @@ containingOpH, [](OpBuilder &b, Location loc) { b.create<transform::ApplyLowerMaskedTransfersPatternsOp>(loc); }); - { - ApplyPatternsOpPatterns configuration; - configuration.rankReducingLinalg = true; - configuration.rankReducingVector = true; - b.create<ApplyPatternsOp>(containingOpH, configuration); - } + b.create<transform::ApplyPatternsOp>( + containingOpH, [](OpBuilder &b, Location loc) { + b.create<transform::ApplyCastAwayVectorLeadingOneDimPatternsOp>(loc); + b.create<transform::ApplyFoldUnitExtentDimsViaSlicesPatternsOp>(loc); + b.create<IREE::transform_dialect:: + ApplyFoldReshapeIntoTensorHalInterfacePatternsOp>(loc); + }); return containingOpH; } @@ -483,8 +484,6 @@ Value mlir::iree_compiler::buildMemoryOptimizations(ImplicitLocOpBuilder &b, Value funcH) { - ApplyPatternsOpPatterns configuration; - configuration.rankReducingVector = true; // Apply canonicalizations and enablings twice as they enable each other. for (int i = 0; i < 2; ++i) { buildCanonicalizationAndEnablingTransforms(
diff --git a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/Common.cpp b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/Common.cpp index 47d7130..ba43176 100644 --- a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/Common.cpp +++ b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/Common.cpp
@@ -153,10 +153,10 @@ Value variantH, Value funcH, int64_t warpSize) { - ApplyPatternsOpPatterns patterns; - patterns.foldMemrefAliases = true; - patterns.rankReducingVector = true; - b.create<ApplyPatternsOp>(funcH, patterns); + b.create<transform::ApplyPatternsOp>(funcH, [](OpBuilder &b, Location loc) { + b.create<transform::ApplyFoldMemrefAliasOpsPatternsOp>(loc); + b.create<transform::ApplyCastAwayVectorLeadingOneDimPatternsOp>(loc); + }); Value ifH = b.create<MatchOp>(funcH, scf::IfOp::getOperationName()); // Locally suppress failures for this op only because it doesn't cover the // `threadIdx.x == 0 && threadIdx.y == 0` case at the moment. @@ -530,11 +530,9 @@ // TODO: Fewer canonicalization. iree_compiler::buildCanonicalizationAndEnablingTransforms(b, funcH); b.create<iree_compiler::IREE::transform_dialect::HoistStaticAllocOp>(funcH); - { - ApplyPatternsOpPatterns config; - config.foldMemrefAliases = true; - b.create<ApplyPatternsOp>(funcH, config); - } + b.create<transform::ApplyPatternsOp>(funcH, [](OpBuilder &b, Location loc) { + b.create<transform::ApplyFoldMemrefAliasOpsPatternsOp>(loc); + }); { ApplyPatternsOpPatterns config; config.extractAddressComputations = true; @@ -569,9 +567,9 @@ ImplicitLocOpBuilder &b, Value funcH, const AbstractGemmLikeStrategy &strategy) { iree_compiler::buildCanonicalizationAndEnablingTransforms(b, funcH); - ApplyPatternsOpPatterns config; - config.foldMemrefAliases = true; - b.create<ApplyPatternsOp>(funcH, config); + b.create<transform::ApplyPatternsOp>(funcH, [](OpBuilder &b, Location loc) { + b.create<transform::ApplyFoldMemrefAliasOpsPatternsOp>(loc); + }); // TODO: Avoid brittle matching here. // TODO: Better builder after integrate. Value allocH = b.create<transform::MatchOp>(
diff --git a/tests/transform_dialect/cpu/attention_codegen_spec.mlir b/tests/transform_dialect/cpu/attention_codegen_spec.mlir index a91ef15..a541e66 100644 --- a/tests/transform_dialect/cpu/attention_codegen_spec.mlir +++ b/tests/transform_dialect/cpu/attention_codegen_spec.mlir
@@ -24,7 +24,11 @@ // Vectorize function // ========================================== %func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %func_3 = transform.structured.vectorize %func : (!transform.any_op) -> !transform.any_op transform.apply_patterns to %func_3 { transform.apply_patterns.iree.fold_fill_into_pad
diff --git a/tests/transform_dialect/cuda/double_mma_layout_analysis_codegen_spec.mlir b/tests/transform_dialect/cuda/double_mma_layout_analysis_codegen_spec.mlir index a2d9a75..09e3cc0 100644 --- a/tests/transform_dialect/cuda/double_mma_layout_analysis_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/double_mma_layout_analysis_codegen_spec.mlir
@@ -23,7 +23,11 @@ // Step 3. Vectorize // =========================================================================== %func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %func_3 = transform.structured.vectorize %func : (!transform.any_op) -> !transform.any_op // Step 4. Bufferize
diff --git a/tests/transform_dialect/cuda/eltwise_reduction_codegen_spec.mlir b/tests/transform_dialect/cuda/eltwise_reduction_codegen_spec.mlir index 83f4f57..d9beb7f 100644 --- a/tests/transform_dialect/cuda/eltwise_reduction_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/eltwise_reduction_codegen_spec.mlir
@@ -63,7 +63,11 @@ // Step 4. Rank-reduce and vectorize. // =========================================================================== %func_1 = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func_1 { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func_1 { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %func_3 = transform.structured.vectorize %func_1 : (!transform.any_op) -> !transform.any_op // Step 5. Bufferize and drop HAL decriptor from memref ops. @@ -81,7 +85,11 @@ // Step 7. Post-bufferization vector distribution with rank-reduction. // =========================================================================== - transform.iree.apply_patterns %func_4 { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func_4 { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %if_op = transform.structured.match ops{["scf.if"]} in %variant_op_2 : (!transform.any_op) -> !transform.any_op // Don't complain about unsupported if (threadIdx.x == 0 && threadIdx.y == 0) // at this point.
diff --git a/tests/transform_dialect/cuda/eltwise_reduction_eltwise_codegen_spec.mlir b/tests/transform_dialect/cuda/eltwise_reduction_eltwise_codegen_spec.mlir index fc92816..bb2598b 100644 --- a/tests/transform_dialect/cuda/eltwise_reduction_eltwise_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/eltwise_reduction_eltwise_codegen_spec.mlir
@@ -70,7 +70,11 @@ // Step 4. Rank-reduce and vectorize. // =========================================================================== %func_1 = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func_1 { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func_1 { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %func_2 = transform.structured.vectorize %func_1 : (!transform.any_op) -> !transform.any_op // Step 5. Bufferize and drop HAL decriptor from memref ops. @@ -88,7 +92,12 @@ // Step 7. Post-bufferization vector distribution with rank-reduction. // =========================================================================== - transform.iree.apply_patterns %func_3 { rank_reducing_linalg, rank_reducing_vector, fold_memref_aliases } : (!transform.any_op) -> () + transform.apply_patterns to %func_3 { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.memref.fold_memref_alias_ops + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %if_op = transform.structured.match ops{["scf.if"]} in %variant_op_2 : (!transform.any_op) -> !transform.any_op // Don't complain about unsupported if (threadIdx.x == 0 && threadIdx.y == 0) // at this point.
diff --git a/tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir b/tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir index 04777a4..5d06e43 100644 --- a/tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/mma_elemwise_layout_analysis_codegen_spec.mlir
@@ -21,7 +21,11 @@ // Step 3. Vectorize // =========================================================================== %func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %func_3 = transform.structured.vectorize %func : (!transform.any_op) -> !transform.any_op // Step 4. Bufferize
diff --git a/tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir b/tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir index 957cc8d..84a6ce0 100644 --- a/tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/mma_reduction_layout_analysis_codegen_spec.mlir
@@ -22,7 +22,11 @@ // Step 3. Vectorize // =========================================================================== %func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %func_3 = transform.structured.vectorize %func : (!transform.any_op) -> !transform.any_op // Step 4. Bufferize
diff --git a/tests/transform_dialect/cuda/mma_using_layout_analysis_codegen_spec.mlir b/tests/transform_dialect/cuda/mma_using_layout_analysis_codegen_spec.mlir index 25914a8..c647bca 100644 --- a/tests/transform_dialect/cuda/mma_using_layout_analysis_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/mma_using_layout_analysis_codegen_spec.mlir
@@ -26,7 +26,11 @@ // Step 3. Vectorize // =========================================================================== %func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %func_3 = transform.structured.vectorize %func : (!transform.any_op) -> !transform.any_op // Step 4. Bufferize
diff --git a/tests/transform_dialect/cuda/reduction_codegen_spec.mlir b/tests/transform_dialect/cuda/reduction_codegen_spec.mlir index e54b8d8..6902576 100644 --- a/tests/transform_dialect/cuda/reduction_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/reduction_codegen_spec.mlir
@@ -56,7 +56,11 @@ // =========================================================================== %func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %func_3 = transform.structured.vectorize %func : (!transform.any_op) -> !transform.any_op // Step 5. Bufferize and drop HAL decriptor from memref ops. @@ -78,7 +82,12 @@ // Step 7. Post-bufferization vector distribution with rank-reduction. // =========================================================================== - transform.iree.apply_patterns %func_5 { rank_reducing_linalg, rank_reducing_vector, fold_memref_aliases } : (!transform.any_op) -> () + transform.apply_patterns to %func_5 { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.memref.fold_memref_alias_ops + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %if_op = transform.structured.match ops{["scf.if"]} in %variant_op_3 : (!transform.any_op) -> !transform.any_op // Don't complain about unsupported if (threadIdx.x == 0 && threadIdx.y == 0)
diff --git a/tests/transform_dialect/cuda/reduction_eltwise_codegen_spec.mlir b/tests/transform_dialect/cuda/reduction_eltwise_codegen_spec.mlir index ff4a48b..3bd414b 100644 --- a/tests/transform_dialect/cuda/reduction_eltwise_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/reduction_eltwise_codegen_spec.mlir
@@ -91,7 +91,11 @@ // Step 4. Rank-reduce and vectorize. // =========================================================================== %func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %func_3 = transform.structured.vectorize %func : (!transform.any_op) -> !transform.any_op // Step 5. Bufferize and drop HAL decriptor from memref ops. @@ -112,7 +116,12 @@ // Step 7. Post-bufferization vector distribution with rank-reduction. // =========================================================================== - transform.iree.apply_patterns %func_5 { rank_reducing_linalg, rank_reducing_vector, fold_memref_aliases } : (!transform.any_op) -> () + transform.apply_patterns to %func_5 { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.memref.fold_memref_alias_ops + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %if_op = transform.structured.match ops{["scf.if"]} in %variant_op_3 : (!transform.any_op) -> !transform.any_op // Don't complain about unsupported if (threadIdx.x == 0 && threadIdx.y == 0)
diff --git a/tests/transform_dialect/cuda/reduction_v2_codegen_spec.mlir b/tests/transform_dialect/cuda/reduction_v2_codegen_spec.mlir index e8fa225..00958bb 100644 --- a/tests/transform_dialect/cuda/reduction_v2_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/reduction_v2_codegen_spec.mlir
@@ -41,7 +41,11 @@ // Step 4. Rank-reduce and vectorize. // =========================================================================== %func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %func_3 = transform.structured.vectorize %func : (!transform.any_op) -> !transform.any_op // Step 5. Bufferize and drop HAL decriptor from memref ops. @@ -71,7 +75,12 @@ // Step 7. Post-bufferization vector distribution with rank-reduction. // =========================================================================== - transform.iree.apply_patterns %func_7 { rank_reducing_linalg, rank_reducing_vector, fold_memref_aliases } : (!transform.any_op) -> () + transform.apply_patterns to %func_7 { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.memref.fold_memref_alias_ops + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %if_op = transform.structured.match ops{["scf.if"]} in %variant_op_3 : (!transform.any_op) -> !transform.any_op %warp = transform.iree.vector.to_warp_execute_on_lane_0 %if_op { warp_size = 32 } : (!transform.any_op) -> !transform.any_op
diff --git a/tests/transform_dialect/cuda/reduction_v3_codegen_spec.mlir b/tests/transform_dialect/cuda/reduction_v3_codegen_spec.mlir index 1740b10..588c089 100644 --- a/tests/transform_dialect/cuda/reduction_v3_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/reduction_v3_codegen_spec.mlir
@@ -60,8 +60,11 @@ : (!transform.any_op) -> !transform.any_op // TODO: masked vectorization on block_more_parallel_op_2 if we want // vector<4> to work as intended. - transform.iree.apply_patterns %func - { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %func_3 = transform.structured.vectorize %func : (!transform.any_op) -> !transform.any_op // Canonicalizations is necessary to get rid of some tensor.cast that block @@ -108,7 +111,12 @@ // Step 6. Post-bufferization vector distribution with rank-reduction. // =========================================================================== - transform.iree.apply_patterns %func_m { rank_reducing_linalg, rank_reducing_vector, fold_memref_aliases } : (!transform.any_op) -> () + transform.apply_patterns to %func_m { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.memref.fold_memref_alias_ops + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %if_op = transform.structured.match ops{["scf.if"]} in %variant_op_3 : (!transform.any_op) -> !transform.any_op %warp = transform.iree.vector.to_warp_execute_on_lane_0 %if_op { warp_size = 32 } : (!transform.any_op) -> !transform.any_op
diff --git a/tests/transform_dialect/cuda/softmax_codegen_spec.mlir b/tests/transform_dialect/cuda/softmax_codegen_spec.mlir index e1fad5a..52e667c 100644 --- a/tests/transform_dialect/cuda/softmax_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/softmax_codegen_spec.mlir
@@ -71,7 +71,11 @@ // Step 3. Rank-reduce and vectorize. // ================================== %func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op transform.structured.vectorize %func : (!transform.any_op) -> !transform.any_op // Step 4. Bufferize and drop HAL decriptor from memref ops. @@ -90,7 +94,12 @@ // Step 6. Post-bufferization vector distribution with rank-reduction. // =================================================================== %end_func = transform.structured.match ops{["func.func"]} in %variant_op_3 : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %end_func { rank_reducing_linalg, rank_reducing_vector, fold_memref_aliases } : (!transform.any_op) -> () + transform.apply_patterns to %end_func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.memref.fold_memref_alias_ops + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %if_op = transform.structured.match ops{["scf.if"]} in %variant_op_3 : (!transform.any_op) -> !transform.any_op %warp = transform.iree.vector.to_warp_execute_on_lane_0 %if_op { warp_size = 32 } : (!transform.any_op) -> !transform.any_op transform.iree.vector.warp_distribute %end_func : (!transform.any_op) -> ()
diff --git a/tests/transform_dialect/cuda/softmax_partial_codegen_spec.mlir b/tests/transform_dialect/cuda/softmax_partial_codegen_spec.mlir index 65b6e2f..ecf4ff7 100644 --- a/tests/transform_dialect/cuda/softmax_partial_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/softmax_partial_codegen_spec.mlir
@@ -58,7 +58,11 @@ // Step 3. Rank-reduce and vectorize. // ================================== %func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op transform.structured.vectorize %func : (!transform.any_op) -> !transform.any_op // Step 4. Bufferize and drop HAL decriptor from memref ops. @@ -78,7 +82,12 @@ // Step 6. Post-bufferization vector distribution with rank-reduction. // =================================================================== %end_func = transform.structured.match ops{["func.func"]} in %variant_op_3 : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %end_func { rank_reducing_linalg, rank_reducing_vector, fold_memref_aliases } : (!transform.any_op) -> () + transform.apply_patterns to %end_func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.memref.fold_memref_alias_ops + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %if_op = transform.structured.match ops{["scf.if"]} in %variant_op_3 : (!transform.any_op) -> !transform.any_op %warp = transform.iree.vector.to_warp_execute_on_lane_0 %if_op { warp_size = 32 } : (!transform.any_op) -> !transform.any_op
diff --git a/tests/transform_dialect/cuda/softmax_v2_codegen_spec.mlir b/tests/transform_dialect/cuda/softmax_v2_codegen_spec.mlir index 93115ab..a062ebf 100644 --- a/tests/transform_dialect/cuda/softmax_v2_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/softmax_v2_codegen_spec.mlir
@@ -86,7 +86,11 @@ // Step 3. Rank-reduce and vectorize. // ================================== %funcx_2 = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %funcx_2 { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %funcx_2 { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op transform.structured.vectorize %funcx_2 : (!transform.any_op) -> !transform.any_op // Step 4. Bufferize and drop HAL decriptor from memref ops. @@ -104,7 +108,12 @@ // Step 6. Post-bufferization vector distribution with rank-reduction. // =================================================================== - transform.iree.apply_patterns %memref_func { rank_reducing_linalg, rank_reducing_vector, fold_memref_aliases } : (!transform.any_op) -> () + transform.apply_patterns to %memref_func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.memref.fold_memref_alias_ops + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op %if_op = transform.structured.match ops{["scf.if"]} in %variant_op_3 : (!transform.any_op) -> !transform.any_op %warp = transform.iree.vector.to_warp_execute_on_lane_0 %if_op { warp_size = 32 }
diff --git a/tests/transform_dialect/cuda/vecadd2d_codegen_spec.mlir b/tests/transform_dialect/cuda/vecadd2d_codegen_spec.mlir index 8a2002e..334ceda 100644 --- a/tests/transform_dialect/cuda/vecadd2d_codegen_spec.mlir +++ b/tests/transform_dialect/cuda/vecadd2d_codegen_spec.mlir
@@ -12,7 +12,11 @@ // Step 2. Rank reduce and bufferize and drop HAL decriptor from memref ops. // =========================================================================== %func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op - transform.iree.apply_patterns %func { rank_reducing_linalg, rank_reducing_vector } : (!transform.any_op) -> () + transform.apply_patterns to %func { + transform.apply_patterns.iree.fold_reshape_into_tensor_hal_interface + transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices + transform.apply_patterns.vector.cast_away_vector_leading_one_dim + } : !transform.any_op transform.iree.eliminate_empty_tensors %variant_op : (!transform.any_op) -> () %variant_op_3 = transform.iree.bufferize { target_gpu } %variant_op : (!transform.any_op) -> !transform.any_op %memref_func = transform.structured.match ops{["func.func"]} in %variant_op_3 : (!transform.any_op) -> !transform.any_op