[mlir] python bindings for vector transform ops

[ftynse]

Summary

Provide Python bindings for transform ops defined in the vector dialect.
All of these ops are sufficiently simple that no mixins are necessary
for them to be nicely usable.

[ingomueller-net]

LGTM! Thanks for showing how to use the enums and make nicer tests -- good input for for future patches!

mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td

let description = [{
 
     This is usally a late step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
   }];
 
   let assemblyFormat = "attr-dict";
 }
 
-// TODO: evolve lowering_strategy to proper enums.
 def ApplyLowerContractionPatternsOp : Op<Transform_Dialect,
     "apply_patterns.vector.lower_contraction",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
   let description = [{
     Indicates that vector contraction-like operations should be lowered to
     finer-grained vector primitives.
 
     This is usually a late step that is run after bufferization as part of the
let description = [{
 
     This is usually the last step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
   }];
 
   let assemblyFormat = "attr-dict";
 }
 
-// TODO: evolve lowering_strategy to proper enums.
 def ApplyLowerMultiReductionPatternsOp : Op<Transform_Dialect,
     "apply_patterns.vector.lower_multi_reduction",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
   let description = [{
     Indicates that vector multi_reduction-like operations should be lowered to
     finer-grained vector primitives.
 
     This is usually a late step that is run after bufferization as part of the
def ApplyLowerTransferPatternsOp : Op<Transform_Dialect,
 
   let arguments = (ins DefaultValuedAttr<I64Attr, "1">:$max_transfer_rank);
 
   let assemblyFormat = [{
     (`max_transfer_rank` `=` $max_transfer_rank^)? attr-dict
   }];
 }
 
-// TODO: evolve lowering_strategy to proper enums.
 def ApplyLowerTransposePatternsOp : Op<Transform_Dialect,
     "apply_patterns.vector.lower_transpose",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
   let description = [{
     Indicates that vector transpose-like operations should be lowered to
     finer-grained vector primitives.
 
     This is usually a late step that is run after bufferization as part of the
let assemblyFormat = [{
     oilist (
       `lowering_strategy` `=` $lowering_strategy
       | `avx2_lowering_strategy` `=` $avx2_lowering_strategy
     )
     attr-dict
   }];
 }
 
-// TODO: evolve split_transfer_strategy to proper enums.
 def ApplySplitTransferFullPartialPatternsOp : Op<Transform_Dialect,
     "apply_patterns.vector.split_transfer_full_partial",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
   let description = [{
     Indicates that vector transfer operations should be split to full and
     partial parts.
 
     This is usually a late step that is run after bufferization as part of the

mlir/python/CMakeLists.txt

declare_mlir_dialect_extension_python_bindings(
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/LinalgStructuredTransformOps.td
   SOURCES
     dialects/_structured_transform_ops_ext.py
     dialects/transform/structured.py
   DIALECT_NAME transform
   EXTENSION_NAME structured_transform)
 
+declare_mlir_dialect_extension_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/VectorTransformOps.td
+  SOURCES
+    dialects/transform/vector.py
+  DIALECT_NAME transform
+  EXTENSION_NAME vector_transform)
+
+set(LLVM_TARGET_DEFINITIONS "${CMAKE_CURRENT_SOURCE_DIR}/mlir/dialects/VectorTransformOps.td")
+mlir_tablegen("dialects/_vector_transform_enum_gen.py" -gen-python-enum-bindings)
+add_public_tablegen_target(MLIRVectorTransformPyEnumGen)
+declare_mlir_python_sources(
+  MLIRPythonSources.Dialects.vector_transform.enum_gen
+  ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
+  ADD_TO_PARENT MLIRPythonSources.Dialects.vector_transform
+  SOURCES "dialects/_vector_transform_enum_gen.py" )
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/MathOps.td
   SOURCES dialects/math.py
   DIALECT_NAME math)
 
 declare_mlir_dialect_python_bindings(

mlir/python/mlir/dialects/VectorTransformOps.td

+//===-- VectorTransformOps.td ------------------------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Entry point of the Python bindings generator for the vector transform ops.
+//
+//===----------------------------------------------------------------------===//
+
+
+#ifndef PYTHON_BINDINGS_VECTORTRANSFORMOPS
+#define PYTHON_BINDINGS_VECTORTRANSFORMOPS
+
+include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.td"
+
+#endif // PYTHON_BINDINGS_VECTORTRANSFORMOPS

mlir/python/mlir/dialects/transform/vector.py

+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from .._vector_transform_enum_gen import *
+from .._vector_transform_ops_gen import *

mlir/test/python/dialects/transform_vector_ext.py

+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import transform
+from mlir.dialects.transform import vector
+
+
+def run_apply_patterns(f):
+    with Context(), Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            sequence = transform.SequenceOp(
+                transform.FailurePropagationMode.PROPAGATE,
+                [],
+                transform.AnyOpType.get(),
+            )
+            with InsertionPoint(sequence.body):
+                apply = transform.ApplyPatternsOp(sequence.bodyTarget)
+                with InsertionPoint(apply.patterns):
+                    f()
+                transform.YieldOp()
+        print("\nTEST:", f.__name__)
+        print(module)
+    return f
+
+
+@run_apply_patterns
+def non_configurable_patterns():
+    # CHECK-LABEL: TEST: non_configurable_patterns
+    # CHECK: apply_patterns
+    # CHECK: transform.apply_patterns.vector.cast_away_vector_leading_one_dim
+    vector.ApplyCastAwayVectorLeadingOneDimPatternsOp()
+    # CHECK: transform.apply_patterns.vector.rank_reducing_subview_patterns
+    vector.ApplyRankReducingSubviewPatternsOp()
+    # CHECK: transform.apply_patterns.vector.transfer_permutation_patterns
+    vector.ApplyTransferPermutationPatternsOp()
+    # CHECK: transform.apply_patterns.vector.lower_broadcast
+    vector.ApplyLowerBroadcastPatternsOp()
+    # CHECK: transform.apply_patterns.vector.lower_masks
+    vector.ApplyLowerMasksPatternsOp()
+    # CHECK: transform.apply_patterns.vector.lower_masked_transfers
+    vector.ApplyLowerMaskedTransfersPatternsOp()
+    # CHECK: transform.apply_patterns.vector.materialize_masks
+    vector.ApplyMaterializeMasksPatternsOp()
+    # CHECK: transform.apply_patterns.vector.lower_outerproduct
+    vector.ApplyLowerOuterProductPatternsOp()
+    # CHECK: transform.apply_patterns.vector.lower_gather
+    vector.ApplyLowerGatherPatternsOp()
+    # CHECK: transform.apply_patterns.vector.lower_scan
+    vector.ApplyLowerScanPatternsOp()
+    # CHECK: transform.apply_patterns.vector.lower_shape_cast
+    vector.ApplyLowerShapeCastPatternsOp()
+
+
+@run_apply_patterns
+def configurable_patterns():
+    # CHECK-LABEL: TEST: configurable_patterns
+    # CHECK: apply_patterns
+    # CHECK: transform.apply_patterns.vector.lower_transfer
+    # CHECK-SAME: max_transfer_rank = 4
+    vector.ApplyLowerTransferPatternsOp(max_transfer_rank=4)
+    # CHECK: transform.apply_patterns.vector.transfer_to_scf
+    # CHECK-SAME: max_transfer_rank = 3
+    # CHECK-SAME: full_unroll = true
+    vector.ApplyTransferToScfPatternsOp(max_transfer_rank=3, full_unroll=True)
+
+
+@run_apply_patterns
+def enum_configurable_patterns():
+    # CHECK: transform.apply_patterns.vector.lower_contraction
+    vector.ApplyLowerContractionPatternsOp()
+    # CHECK: transform.apply_patterns.vector.lower_contraction
+    # CHECK-SAME: lowering_strategy = matmulintrinsics
+    vector.ApplyLowerContractionPatternsOp(
+        lowering_strategy=vector.VectorContractLowering.MATMUL
+    )
+    # CHECK: transform.apply_patterns.vector.lower_contraction
+    # CHECK-SAME: lowering_strategy = parallelarith
+    vector.ApplyLowerContractionPatternsOp(
+        lowering_strategy=vector.VectorContractLowering.PARALLEL_ARITH
+    )
+
+    # CHECK: transform.apply_patterns.vector.lower_multi_reduction
+    vector.ApplyLowerMultiReductionPatternsOp()
+    # CHECK: transform.apply_patterns.vector.lower_multi_reduction
+    # This is the default mode, not printed.
+    vector.ApplyLowerMultiReductionPatternsOp(
+        lowering_strategy=vector.VectorMultiReductionLowering.INNER_PARALLEL
+    )
+    # CHECK: transform.apply_patterns.vector.lower_multi_reduction
+    # CHECK-SAME: lowering_strategy = innerreduction
+    vector.ApplyLowerMultiReductionPatternsOp(
+        lowering_strategy=vector.VectorMultiReductionLowering.INNER_REDUCTION
+    )
+
+    # CHECK: transform.apply_patterns.vector.lower_transpose
+    # CHECK-SAME: lowering_strategy = eltwise
+    # CHECK-SAME: avx2_lowering_strategy = false
+    vector.ApplyLowerTransposePatternsOp()
+    # CHECK: transform.apply_patterns.vector.lower_transpose
+    # CHECK-SAME: lowering_strategy = eltwise
+    # CHECK-SAME: avx2_lowering_strategy = false
+    vector.ApplyLowerTransposePatternsOp(
+        lowering_strategy=vector.VectorTransposeLowering.ELT_WISE
+    )
+    # CHECK: transform.apply_patterns.vector.lower_transpose
+    # CHECK-SAME: lowering_strategy = flat_transpose
+    # CHECK-SAME: avx2_lowering_strategy = false
+    vector.ApplyLowerTransposePatternsOp(
+        lowering_strategy=vector.VectorTransposeLowering.FLAT
+    )
+    # CHECK: transform.apply_patterns.vector.lower_transpose
+    # CHECK-SAME: lowering_strategy = shuffle_1d
+    # CHECK-SAME: avx2_lowering_strategy = false
+    vector.ApplyLowerTransposePatternsOp(
+        lowering_strategy=vector.VectorTransposeLowering.SHUFFLE1_D
+    )
+    # CHECK: transform.apply_patterns.vector.lower_transpose
+    # CHECK-SAME: lowering_strategy = shuffle_16x16
+    # CHECK-SAME: avx2_lowering_strategy = false
+    vector.ApplyLowerTransposePatternsOp(
+        lowering_strategy=vector.VectorTransposeLowering.SHUFFLE16X16
+    )
+    # CHECK: transform.apply_patterns.vector.lower_transpose
+    # CHECK-SAME: lowering_strategy = flat_transpose
+    # CHECK-SAME: avx2_lowering_strategy = true
+    vector.ApplyLowerTransposePatternsOp(
+        lowering_strategy=vector.VectorTransposeLowering.FLAT,
+        avx2_lowering_strategy=True,
+    )
+
+    # CHECK: transform.apply_patterns.vector.split_transfer_full_partial
+    vector.ApplySplitTransferFullPartialPatternsOp()
+    # CHECK: transform.apply_patterns.vector.split_transfer_full_partial
+    # CHECK-SAME: split_transfer_strategy = none
+    vector.ApplySplitTransferFullPartialPatternsOp(
+        split_transfer_strategy=vector.VectorTransferSplit.NONE
+    )
+    # CHECK: transform.apply_patterns.vector.split_transfer_full_partial
+    # CHECK-SAME: split_transfer_strategy = "vector-transfer"
+    vector.ApplySplitTransferFullPartialPatternsOp(
+        split_transfer_strategy=vector.VectorTransferSplit.VECTOR_TRANSFER
+    )
+    # CHECK: transform.apply_patterns.vector.split_transfer_full_partial
+    # This is the default mode, not printed.
+    vector.ApplySplitTransferFullPartialPatternsOp(
+        split_transfer_strategy=vector.VectorTransferSplit.LINALG_COPY
+    )
+    # CHECK: transform.apply_patterns.vector.split_transfer_full_partial
+    # CHECK-SAME: split_transfer_strategy = "force-in-bounds"
+    vector.ApplySplitTransferFullPartialPatternsOp(
+        split_transfer_strategy=vector.VectorTransferSplit.FORCE_IN_BOUNDS
+    )

utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel

deps = [
         "//mlir:CallInterfacesTdFiles",
         "//mlir:FunctionInterfacesTdFiles",
         "//mlir:OpBaseTdFiles",
         "//mlir:TransformDialectTdFiles",
         "//mlir:TransformPDLExtensionTdFiles",
     ],
 )
 
+gentbl_filegroup(
+    name = "VectorTransformEnumPyGen",
+    tbl_outs = [
+        (
+            ["-gen-python-enum-bindings"],
+            "mlir/dialects/_vector_transform_enum_gen.py",
+        ),
+    ],
+    tblgen = "//mlir:mlir-tblgen",
+    td_file = "mlir/dialects/VectorTransformOps.td",
+    deps = [
+        "//mlir:OpBaseTdFiles",
+        "//mlir:TransformDialectTdFiles",
+        "//mlir:VectorTransformOpsTdFiles",
+    ],
+)
+
+gentbl_filegroup(
+    name = "VectorTransformOpsPyGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-python-op-bindings",
+                "-bind-dialect=transform",
+                "-dialect-extension=vector_transform",
+            ],
+            "mlir/dialects/_vector_transform_ops_gen.py",
+        ),
+    ],
+    tblgen = "//mlir:mlir-tblgen",
+    td_file = "mlir/dialects/VectorTransformOps.td",
+    deps = [
+        "//mlir:OpBaseTdFiles",
+        "//mlir:TransformDialectTdFiles",
+        "//mlir:VectorTransformOpsTdFiles",
+    ],
+)
+
+
 filegroup(
     name = "TransformOpsPyFiles",
     srcs = [
         "mlir/dialects/_bufferization_transform_ops_ext.py",
         "mlir/dialects/_gpu_transform_ops_ext.py",
         "mlir/dialects/_loop_transform_ops_ext.py",
         "mlir/dialects/_structured_transform_ops_ext.py",
         "mlir/dialects/_transform_ops_ext.py",
         "mlir/dialects/_transform_pdl_extension_ops_ext.py",
         ":BufferizationTransformOpsPyGen",
         ":GPUTransformOpsPyGen",
         ":LoopTransformOpsPyGen",
         ":MemRefTransformOpsPyGen",
         ":PDLTransformOpsPyGen",
         ":StructuredTransformOpsPyGen",
         ":TransformEnumPyGen",
         ":TransformOpsPyGen",
+        ":VectorTransformEnumPyGen",
+        ":VectorTransformOpsPyGen",
     ],
 )
 
 filegroup(
     name = "TransformOpsPackagePyFiles",
     srcs = glob(["mlir/dialects/transform/*.py"]),
 )