[CPU] Enable mmt4d ukernels when iree-llvmcpu-enable-ukernels is not …

…set (iree-org#17928)

This is a follow-up for

iree-org@3b5d269.
The previous revision enables mmt4d ukernels only when
`--iree-llvmcpu-enable-ukernels` is passed to IREE tools. If the
`ukernels` attribute is not present in hal.executable.target, it is not
enabled.

The revisions removes the constraint, so users don't need to pass
`--iree-llvmcpu-enable-ukernels=default` to enable the mmt4d ukernels.

Signed-off-by: hanhanW <[email protected]>

compiler/src/iree/compiler/Codegen/Common/CPU/test/lower_to_ukernel_ops.mlir

@@ -29,6 +29,34 @@ func.func @mmt4d_f32f32f32(%arg0 : tensor<?x?x16x1xf32>, %arg1 : tensor<?x?x16x1
 
 // -----
 
+func.func @mmt4d_no_ukernels_attr_f32f32f32(%arg0 : tensor<?x?x16x1xf32>, %arg1 : tensor<?x?x16x1xf32>,
+    %arg2 : tensor<?x?x16x16xf32>) -> tensor<?x?x16x16xf32> attributes {
+  hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f"}>
+} {
+  %0 = linalg.mmt4d ins(%arg0, %arg1 : tensor<?x?x16x1xf32>, tensor<?x?x16x1xf32>)
+      outs(%arg2 : tensor<?x?x16x16xf32>) -> tensor<?x?x16x16xf32>
+  return %0 : tensor<?x?x16x16xf32>
+}
+// CHECK-LABEL: func @mmt4d_no_ukernels_attr_f32f32f32(
+// CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?x16x1xf32>
+// CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?x16x1xf32>
+// CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: tensor<?x?x16x16xf32>
+//  CHECK-DAG:   %[[FLAGS:.+]] = arith.constant {{[0-9]+}} : i32
+//  CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
+//  CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
+//  CHECK-DAG:   %[[C1_i32:.+]] = arith.constant 1 : i32
+//  CHECK-DAG:   %[[C16_i32:.+]] = arith.constant 16 : i32
+//  CHECK-DAG:   %[[M:.+]] = tensor.dim %[[ARG0]], %[[C0]]
+//  CHECK-DAG:   %[[N:.+]] = tensor.dim %[[ARG1]], %[[C0]]
+//  CHECK-DAG:   %[[K:.+]] = tensor.dim %[[ARG1]], %[[C1]]
+//      CHECK:   %[[MICRO_KERNEL:.+]]:2 = iree_codegen.ukernel.generic "iree_uk_mmt4d"
+// CHECK-SAME:       ins(%[[ARG0]], %[[ARG1]] :
+// CHECK-SAME:       outs(%[[ARG2]] :
+// CHECK-SAME:       (%[[M]], %[[N]], %[[K]], %[[C16_i32]], %[[C16_i32]], %[[C1_i32]], %[[FLAGS]] :
+//      CHECK:   return %[[MICRO_KERNEL]]#0
+
+// -----
+
 func.func @mmt4d_f32f32f32_with_none_ukernel_enabled(%arg0 : tensor<?x?x16x1xf32>, %arg1 : tensor<?x?x16x1xf32>,
     %arg2 : tensor<?x?x16x16xf32>) -> tensor<?x?x16x16xf32> attributes {
   hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {ukernels = "none", target_triple="x86_64-xyz-xyz", cpu_features="+avx512f"}>

compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp

@@ -104,7 +104,6 @@ void LLVMCPULowerExecutableTargetPass::runOnOperation() {
   pipelineOpts.enableVectorMasking =
       isX86(target) || isRISCV(target) ||
       (isAArch64(target) && hasAnySVEFeature(target));
-  pipelineOpts.enableUkernels = hasUkernel(target);
   pipelineOpts.enableAArch64SSVE =
       isAArch64(target) && hasAnySVEFeature(target) && hasSMEFeature(target);
   pipelineOpts.enableAArch64I8mm = isAArch64(target) && hasI8mmFeature(target);

compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp

@@ -501,11 +501,11 @@ void addMmt4dTilingExpertPassPipeline(OpPassManager &funcPassManager,
 
   funcPassManager.addPass(createLLVMCPUTileAndFusePass(
       static_cast<int64_t>(tilingConfig.getVectorCommonParallelLevel())));
-  if (pipelineOpt.enableUkernels) {
-    funcPassManager.addPass(createCPUPrepareUkernelsPass());
-    funcPassManager.addPass(
-        createCPULowerToUKernelsPass(clSkipIntermediateRoundings));
-  }
+  // The below two passes are nop if the "mmt4d" is explicitly excluded in the
+  // ukernels attribute.
+  funcPassManager.addPass(createCPUPrepareUkernelsPass());
+  funcPassManager.addPass(
+      createCPULowerToUKernelsPass(clSkipIntermediateRoundings));
   funcPassManager.addPass(createLLVMCPUTilePass(
       static_cast<int64_t>(tilingConfig.getVectorReductionLevel())));
 
@@ -545,11 +545,11 @@ void addCPUDataTilingPipeline(OpPassManager &funcPassManager,
                               LLVMCPUPipelineOptions &pipelineOpt) {
   addTileAndDistributePasses(funcPassManager);
 
-  if (pipelineOpt.enableUkernels) {
-    funcPassManager.addPass(createCPUPrepareUkernelsPass());
-    funcPassManager.addPass(
-        createCPULowerToUKernelsPass(clSkipIntermediateRoundings));
-  }
+  // The below two passes are nop if pack/unpack is not specified in ukernels
+  // attribute. By default, they are disabled.
+  funcPassManager.addPass(createCPUPrepareUkernelsPass());
+  funcPassManager.addPass(
+      createCPULowerToUKernelsPass(clSkipIntermediateRoundings));
 
   funcPassManager.addPass(
       createLLVMCPUTilePass(tilingConfig.getVectorCommonParallelLevel()));

compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.h

@@ -149,7 +149,6 @@ struct LLVMCPUPipelineOptions {
   bool enableVectorMasking = false;
   bool enableAArch64SSVE = false;
   bool enableAArch64I8mm = false;
-  bool enableUkernels = false;
   bool lowerToAVX2 = false;
 };
 

compiler/src/iree/compiler/Codegen/Utils/Utils.cpp

@@ -152,10 +152,12 @@ getDefaultEnabledUkernels(IREE::HAL::ExecutableTargetAttr targetAttr) {
 bool hasUkernel(IREE::HAL::ExecutableTargetAttr targetAttr,
                 StringRef ukernelName) {
   auto enabledUkernels = getConfigStringAttr(targetAttr, "ukernels");
-  if (!enabledUkernels) {
-    return false;
+  StringRef enabledUkernelsStr;
+  if (enabledUkernels) {
+    enabledUkernelsStr = enabledUkernels->getValue();
+  } else {
+    enabledUkernelsStr = "default";
   }
-  StringRef enabledUkernelsStr = enabledUkernels->getValue();
   // Resolve `default`.
   if (enabledUkernelsStr == "default") {
     enabledUkernelsStr = getDefaultEnabledUkernels(targetAttr);