Merge branch 'main' into xiaohui/vectorization

Author: Xu, Xiaohui1

.github/actions/setup-mlir-python-bindings/action.yml

@@ -0,0 +1,22 @@
+name: setup-mlir-python-bindings
+description: Fetch and setup MLIR Python bindings
+runs:
+  using: "composite"
+  steps:
+    - name: Set LLVM hash
+      shell: bash
+      run: |
+        echo LLVM_HASH=$(cat cmake/llvm-version.txt) >>$GITHUB_ENV
+
+    - name: Fetch requirements for python binding
+      uses: actions/checkout@v4
+      with:
+        repository: llvm/llvm-project
+        ref: ${{ env.LLVM_HASH }}
+        sparse-checkout: mlir/python/requirements.txt
+        sparse-checkout-cone-mode: false
+        path: llvm-dep
+
+    - name: Install requirements
+      shell: bash
+      run: python3 -m pip install -r llvm-dep/mlir/python/requirements.txt

.github/workflows/build-gpu-imex.yml

@@ -22,21 +22,9 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
-      - name: Set LLVM hash
-        run: |
-          echo LLVM_HASH=$(cat cmake/llvm-version-imex.txt) >>$GITHUB_ENV
-
-      - name: Fetch requirements for python binding
-        uses: actions/checkout@v4
-        with:
-          repository: llvm/llvm-project
-          ref: ${{ env.LLVM_HASH }}
-          sparse-checkout: mlir/python/requirements.txt
-          sparse-checkout-cone-mode: false
-          path: llvm-dep
-
-      - name: Install requirements
-        run: python3 -m pip install -r llvm-dep/mlir/python/requirements.txt
+      - name: Setup MLIR Python bindings
+        id: setup-mlir-python-bindings
+        uses: ./.github/actions/setup-mlir-python-bindings
 
       - name: Build
         run: |

.github/workflows/build-llvm.yml

@@ -32,7 +32,7 @@ jobs:
 
       - uses: actions/checkout@v4
         with:
-          repository: Menooker/mlir-extensions
+          repository: intel/mlir-extensions
           ref: ${{ env.IMEX_HASH }}
           path: mlir-extensions
         if: ${{ matrix.build-type == 'IMEX' }}

.github/workflows/build.yml

@@ -22,21 +22,9 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
-      - name: Set LLVM hash
-        run: |
-          echo LLVM_HASH=$(cat cmake/llvm-version.txt) >>$GITHUB_ENV
-
-      - name: Fetch requirements for python binding
-        uses: actions/checkout@v4
-        with:
-          repository: llvm/llvm-project
-          ref: ${{ env.LLVM_HASH }}
-          sparse-checkout: mlir/python/requirements.txt
-          sparse-checkout-cone-mode: false
-          path: llvm-dep
-
-      - name: Install requirements
-        run: python3 -m pip install -r llvm-dep/mlir/python/requirements.txt
+      - name: Setup MLIR Python bindings
+        id: setup-mlir-python-bindings
+        uses: ./.github/actions/setup-mlir-python-bindings
 
       - name: Build
         run: |

.github/workflows/coverity.yml

@@ -53,6 +53,10 @@ jobs:
           echo "email=$email" | tee -a $GITHUB_ENV
           echo "version=$version" | tee -a $GITHUB_ENV
 
+      - name: Setup MLIR Python bindings
+        id: setup-mlir-python-bindings
+        uses: ./.github/actions/setup-mlir-python-bindings
+
       - name: Run coverity build
         run: |
           

cmake/imex-version.txt

@@ -1 +1 @@
-ee459724294e165e360e1de72ad3b217eb9b6206
+6c2e414a953b9a118bce6adac21cf9d42630e674

cmake/imex.cmake

@@ -14,7 +14,7 @@ if (NOT DEFINED IMEX_INCLUDES)
 
     # TODO: Change to main https://github.com/intel/mlir-extensions when all the
     # required functionality is merged.
-    gc_fetch_content(imex "${IMEX_HASH}" https://github.com/Menooker/mlir-extensions
+    gc_fetch_content(imex "${IMEX_HASH}" https://github.com/intel/mlir-extensions
             SET IMEX_CHECK_LLVM_VERSION=ON IMEX_ENABLE_L0_RUNTIME=0
     )
 

include/gc/Analysis/MatmulConfigAnalysis.h

@@ -109,16 +109,14 @@ getOprandDimType(linalg::LinalgOp &linalgOp) {
   if (llvm::isa<linalg::ContractionOpInterface>(linalgOp.getOperation())) {
     return getContractionOpOperandDimType(linalgOp);
   } else if (linalgx::isGenericPackedMatmulOp(
-                 linalgOp.getOperation(), linalgx::PackingType::VNNI_MM2D) ||
-             llvm::isa<linalgx::Mm2DVnniOp>(linalgOp)) {
+                 linalgOp.getOperation(), linalgx::PackingType::VNNI_MM2D)) {
     return SmallVector<SmallVector<DimType>>{
         SmallVector<DimType>{DimType::M, DimType::K},
         SmallVector<DimType>{DimType::N, DimType::K, DimType::K, DimType::N,
                              DimType::K},
         SmallVector<DimType>{DimType::M, DimType::N, DimType::M, DimType::N}};
   } else if (linalgx::isGenericPackedMatmulOp(
-                 linalgOp.getOperation(), linalgx::PackingType::VNNI_MM4D) ||
-             llvm::isa<linalgx::Mm4DVnniOp>(linalgOp)) {
+                 linalgOp.getOperation(), linalgx::PackingType::VNNI_MM4D)) {
     return SmallVector<SmallVector<DimType>>{
         SmallVector<DimType>{DimType::M, DimType::K, DimType::M, DimType::K},
         SmallVector<DimType>{DimType::N, DimType::K, DimType::K, DimType::N,

include/gc/Dialect/Linalgx/LinalgxStructuredOps.td

@@ -104,212 +104,4 @@ def Linalgx_SigmoidOp : LinalgxStructuredBase_Op<"sigmoid",
   }];
 }
 
-def Linalgx_Mm2DVnniOp 
-    : LinalgxStructuredBase_Op<"mm2d_vnni", [AttrSizedOperandSegments]> {
-  let summary = "Transposed matmul with 2d input and vnni packed weights";
-  let description = [{
-    Supported format: A[M, K] * B[N0, K0, k, n, v] -> C[M, N], with:
-      N = N0 * n
-      K = K0 * k * v; v = (2, 4)
-  }];
-  let arguments = (ins
-      Variadic<TensorOrMemref>:$inputs,
-      Variadic<TensorOrMemref>:$outputs);
-  let results = (outs Variadic<TensorOrMemref>:$results);
-  let regions = (region AnyRegion:$region);
-
-  let skipDefaultBuilders = 1;
-  let builders = [
-      OpBuilder<
-      (ins 
-        "TypeRange":$resultTensorTypes, 
-        "ValueRange":$inputs,
-        "ValueRange":$outputs,
-        CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
-      [{
-        buildStructuredOp($_builder, $_state, resultTensorTypes,
-          inputs, outputs, attributes, Mm2DVnniOp::getRegionBuilder());
-      }]>
-  ];
-
-  let hasCustomAssemblyFormat = 1;
-  let hasFolder = 1;
-  let hasVerifier = 1;
-
-  let extraClassDeclaration = structuredOpsBaseDecls # [{
-    // Declare functions necessary for LinalgStructuredInterface.
-    SmallVector<utils::IteratorType> getIteratorTypesArray();
-    ArrayAttr getIndexingMaps();
-    static unsigned getNumRegionArgs() { return 3; }
-    std::string getLibraryCallName() {
-      return "op_has_no_registered_library_name";
-    }
-
-    // Implement functions necessary for DestinationStyleOpInterface.
-    MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }
-
-    static void regionBuilder(ImplicitLocOpBuilder &b,
-                              Block &block, ArrayRef<NamedAttribute> attrs);
-    static std::function<void(ImplicitLocOpBuilder &,
-                              Block &, ArrayRef<NamedAttribute>)>
-    getRegionBuilder() {
-      return regionBuilder;
-    }
-  }];
-}
-
-def Linalgx_Mm4DVnniOp 
-    : LinalgxStructuredBase_Op<"mm4d_vnni", [AttrSizedOperandSegments]> {
-  let summary = "Transposed matmul with 4d blocking input and vnni packed weights";
-  let description = [{
-    Supported format: A[M, K, m, k] * B[N, K, k0, n, v] -> C[M, N, m, n], with:
-      k = k0 * v; v = (2, 4)
-  }];
-  let arguments = (ins
-      Variadic<TensorOrMemref>:$inputs,
-      Variadic<TensorOrMemref>:$outputs);
-  let results = (outs Variadic<TensorOrMemref>:$results);
-  let regions = (region AnyRegion:$region);
-
-  let skipDefaultBuilders = 1;
-  let builders = [
-      OpBuilder<
-      (ins 
-        "TypeRange":$resultTensorTypes, 
-        "ValueRange":$inputs,
-        "ValueRange":$outputs,
-        CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
-      [{
-        buildStructuredOp($_builder, $_state, resultTensorTypes,
-          inputs, outputs, attributes, Mm4DVnniOp::getRegionBuilder());
-      }]>
-  ];
-
-  let hasCustomAssemblyFormat = 1;
-  let hasFolder = 1;
-  let hasVerifier = 1;
-
-  let extraClassDeclaration = structuredOpsBaseDecls # [{
-    // Declare functions necessary for LinalgStructuredInterface.
-    SmallVector<utils::IteratorType> getIteratorTypesArray();
-    ArrayAttr getIndexingMaps();
-    static unsigned getNumRegionArgs() { return 3; }
-    std::string getLibraryCallName() {
-      return "op_has_no_registered_library_name";
-    }
-
-    // Implement functions necessary for DestinationStyleOpInterface.
-    MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }
-
-    static void regionBuilder(ImplicitLocOpBuilder &b,
-                              Block &block, ArrayRef<NamedAttribute> attrs);
-    static std::function<void(ImplicitLocOpBuilder &,
-                              Block &, ArrayRef<NamedAttribute>)>
-    getRegionBuilder() {
-      return regionBuilder;
-    }
-  }];
-}
-
-def Linalgx_BatchReduceMatmulVnniOp 
-    : LinalgxStructuredBase_Op<"batch_reduce_matmul_vnni", [AttrSizedOperandSegments]> {
-  let summary = "Batch reduced matmul with 3d batch input and vnni packed weights";
-  let description = [{
-    Supported format: A[B, M, K] * B[B, k, N, v] -> C[M, N], with:
-      K = k * v; v = (2, 4)
-  }];
-  let arguments = (ins
-      Variadic<TensorOrMemref>:$inputs,
-      Variadic<TensorOrMemref>:$outputs);
-  let results = (outs Variadic<TensorOrMemref>:$results);
-  let regions = (region AnyRegion:$region);
-
-  let skipDefaultBuilders = 1;
-  let builders = [
-      OpBuilder<
-      (ins 
-        "TypeRange":$resultTensorTypes, 
-        "ValueRange":$inputs,
-        "ValueRange":$outputs,
-        CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
-      [{
-        buildStructuredOp($_builder, $_state, resultTensorTypes,
-          inputs, outputs, attributes, BatchReduceMatmulVnniOp::getRegionBuilder());
-      }]>
-  ];
-
-  let hasCustomAssemblyFormat = 1;
-  let hasFolder = 1;
-  let hasVerifier = 1;
-
-  let extraClassDeclaration = structuredOpsBaseDecls # [{
-    // Declare functions necessary for LinalgStructuredInterface.
-    SmallVector<utils::IteratorType> getIteratorTypesArray();
-    ArrayAttr getIndexingMaps();
-    static unsigned getNumRegionArgs() { return 3; }
-    std::string getLibraryCallName() {
-      return "op_has_no_registered_library_name";
-    }
-
-    // Implement functions necessary for DestinationStyleOpInterface.
-    MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }
-
-    static void regionBuilder(ImplicitLocOpBuilder &b,
-                              Block &block, ArrayRef<NamedAttribute> attrs);
-    static std::function<void(ImplicitLocOpBuilder &,
-                              Block &, ArrayRef<NamedAttribute>)>
-    getRegionBuilder() {
-      return regionBuilder;
-    }
-  }];
-}
-
-def Linalgx_MultiBatchMatmulOp : LinalgxStructuredBase_Op<"multi_batch_matmul", 
-      [AttrSizedOperandSegments, LinalgContractionOpInterface]> {
-  let summary = "Batch matmul with variable batch dims";
-  let arguments = (ins
-      Variadic<TensorOrMemref>:$inputs,
-      Variadic<TensorOrMemref>:$outputs);
-  let results = (outs Variadic<TensorOrMemref>:$results);
-  let regions = (region AnyRegion:$region);
-
-  let skipDefaultBuilders = 1;
-  let builders = [
-      OpBuilder<
-      (ins 
-        "TypeRange":$resultTensorTypes, 
-        "ValueRange":$inputs,
-        "ValueRange":$outputs,
-        CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
-      [{
-        buildStructuredOp($_builder, $_state, resultTensorTypes,
-          inputs, outputs, attributes, MultiBatchMatmulOp::getRegionBuilder());
-      }]>
-  ];
-
-  let hasCustomAssemblyFormat = 1;
-  let hasFolder = 1;
-
-  let extraClassDeclaration = structuredOpsBaseDecls # [{
-    // Declare functions necessary for LinalgStructuredInterface.
-    SmallVector<utils::IteratorType> getIteratorTypesArray();
-    ArrayAttr getIndexingMaps();
-    static unsigned getNumRegionArgs() { return 3; }
-    std::string getLibraryCallName() {
-      return "op_has_no_registered_library_name";
-    }
-
-    // Implement functions necessary for DestinationStyleOpInterface.
-    MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }
-
-    static void regionBuilder(ImplicitLocOpBuilder &b,
-                              Block &block, ArrayRef<NamedAttribute> attrs);
-    static std::function<void(ImplicitLocOpBuilder &,
-                              Block &, ArrayRef<NamedAttribute>)>
-    getRegionBuilder() {
-      return regionBuilder;
-    }
-  }];
-}
-
 #endif // LINALGX_STRUCTURED_OPS

include/gc/Dialect/Linalgx/Utils.h

@@ -18,7 +18,7 @@ namespace mlir {
 namespace linalgx {
 
 /// @brief enum of type of matmul packing
-enum class PackingType {
+enum class PackingType : int {
   MM4D = 0,    // MKmk x NKkn
   VNNI_MM2D,   // MK x NKknV
   VNNI_MM4D,   // MKmk x NKknV
@@ -43,6 +43,30 @@ makeGenericPackedMatmulOp(OpBuilder &builder, Location loc, PackingType opType,
 /// @return true if op is a generic packed matmul Op
 bool isGenericPackedMatmulOp(Operation *op, PackingType opType);
 
+template <typename... Args>
+inline bool isGenericPackedMatmulOp(Operation *op, PackingType first,
+                                    Args... args) {
+  return isGenericPackedMatmulOp(op, first) ||
+         isGenericPackedMatmulOp(op, args...);
+}
+
+/// @brief identify a generic packed matmul Op based on any PackingType
+/// @param op the op
+/// @return true if op is a generic packed matmul Op
+template <int T, int N> inline bool isAnyGenericPackedMatmulOp(Operation *op) {
+  return isGenericPackedMatmulOp(op, (PackingType)N) ||
+         isAnyGenericPackedMatmulOp<T + 1, N>(op);
+}
+constexpr int NUM_ALL_TYPES = (int)PackingType::NUM_TYPES;
+template <>
+inline bool
+isAnyGenericPackedMatmulOp<NUM_ALL_TYPES, NUM_ALL_TYPES>(Operation *op) {
+  return false;
+}
+inline bool isAnyGenericPackedMatmulOp(Operation *op) {
+  return isAnyGenericPackedMatmulOp<0, NUM_ALL_TYPES>(op);
+}
+
 /// @brief identify a matmul Op based on ContractionOp and PackingType
 /// @param op the op
 /// @return true if op is a matmul Op