Add nop pass to different backend.

Signed-off-by: Alan Li <me@alanli.org>
nod-ai · Jul 13, 2024 · 2ed3f92 · 2ed3f92
1 parent be461bd
commit 2ed3f92
Show file tree

Hide file tree

Showing 6 changed files with 254 additions and 2 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoNop.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoNop.cpp
@@ -6,12 +6,15 @@
 
 #include "iree/compiler/Codegen/Common/EncodingUtils.h"
 #include "iree/compiler/Codegen/Common/PassDetail.h"
+#include "iree/compiler/Codegen/Common/PassUtils.h"
 #include "iree/compiler/Codegen/Common/Passes.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/Passes.h"
 
 namespace mlir::iree_compiler {
 
@@ -84,4 +87,10 @@ createMaterializeEncodingIntoNopPass() {
   return std::make_unique<MaterializeEncodingIntoNopPass>();
 }
 
+void addEncodingToNopPasses(FunctionLikeNest &passManager) {
+  passManager.addPass(createMaterializeEncodingIntoNopPass)
+      .addPass(createBufferizeCopyOnlyDispatchesPass)
+      .addPass(createCanonicalizerPass);
+}
+
 } // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/Common/Passes.h b/compiler/src/iree/compiler/Codegen/Common/Passes.h
@@ -337,6 +337,9 @@ void populateVectorTransferTensorSliceTransforms(RewritePatternSet &patterns,
 /// Method to register all passes.
 void registerCodegenCommonPasses();
 
+/// Populate Encoding to Nop pass and canonicalizer pass to the pipeline
+void addEncodingToNopPasses(FunctionLikeNest &passManager);
+
 } // namespace mlir::iree_compiler
 
 #endif // IREE_COMPILER_CODEGEN_COMMON_PASSES_H_
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
@@ -1040,9 +1040,9 @@ static void buildLLVMGPUCodegenConfigurationPassPipelineImpl(
     FunctionLikeNest funcPassManager(modulePassManager);
     funcPassManager.addPass(createGPUGeneralizeNamedOpsPass);
     addCommonTargetExecutablePreprocessingPasses(funcPassManager);
+    addEncodingToNopPasses(funcPassManager);
   }
   modulePassManager.addPass(createMaterializeUserConfigsPass());
-
   modulePassManager.addPass(createLLVMGPUSelectLoweringStrategyPass());
 }
 

diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/Passes.cpp b/compiler/src/iree/compiler/Codegen/SPIRV/Passes.cpp
@@ -623,9 +623,9 @@ static void buildSPIRVCodegenConfigurationPassPipelineImpl(
     FunctionLikeNest funcPassManager(modulePassManager);
     funcPassManager.addPass(createGPUGeneralizeNamedOpsPass);
     addCommonTargetExecutablePreprocessingPasses(funcPassManager);
+    addEncodingToNopPasses(funcPassManager);
   }
   modulePassManager.addPass(createMaterializeUserConfigsPass());
-
   modulePassManager.addPass(createSPIRVSelectLoweringStrategyPass());
 }
 

diff --git a/tests/e2e/matmul/BUILD.bazel b/tests/e2e/matmul/BUILD.bazel
@@ -288,6 +288,91 @@ X86_64_AVX512_BF16 = X86_64_AVX512 + [
     "large",
 ]]
 
+[iree_generated_e2e_runner_test(
+    name = "e2e_matmul_vmvx_experimental_dt%s_%s_%s" % (
+        ("_uk" if use_uk else ""),
+        lhs_rhs_type,
+        acc_type,
+    ),
+    compiler_flags = [
+        "--iree-opt-data-tiling",
+        "--iree-global-opt-enable-early-materialization=false",
+    ],
+    generator = ":generate_e2e_matmul_tests",
+    generator_args = [
+        "--lhs_rhs_type=%s" % lhs_rhs_type,
+        "--acc_type=%s" % acc_type,
+        "--shapes=small",
+    ],
+    tags = [],
+    target_backends_and_drivers = [
+        ("vmvx", "local-task"),
+    ],
+    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
+    test_type = "matmul",
+) for use_uk in [
+    False,
+    True,
+] for (lhs_rhs_type, acc_type) in (
+    [
+        ("f32", "f32"),
+    ]
+)]
+
+[iree_generated_e2e_runner_test(
+    name = "e2e_matmul_cuda_experimental_dt_%s_%s" % (
+        lhs_rhs_type,
+        acc_type,
+    ),
+    compiler_flags = [
+        "--iree-opt-data-tiling",
+        "--iree-global-opt-enable-early-materialization=false",
+    ],
+    generator = ":generate_e2e_matmul_tests",
+    generator_args = [
+        "--lhs_rhs_type=%s" % lhs_rhs_type,
+        "--acc_type=%s" % acc_type,
+        "--shapes=small",
+    ],
+    tags = [],
+    target_backends_and_drivers = [
+        ("cuda", "cuda"),
+    ],
+    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
+    test_type = "matmul",
+) for (lhs_rhs_type, acc_type) in (
+    [
+        ("f32", "f32"),
+    ]
+)]
+
+[iree_generated_e2e_runner_test(
+    name = "e2e_matmul_spirv_experimental_dt_%s_%s" % (
+        lhs_rhs_type,
+        acc_type,
+    ),
+    compiler_flags = [
+        "--iree-opt-data-tiling",
+        "--iree-global-opt-enable-early-materialization=false",
+    ],
+    generator = ":generate_e2e_matmul_tests",
+    generator_args = [
+        "--lhs_rhs_type=%s" % lhs_rhs_type,
+        "--acc_type=%s" % acc_type,
+        "--shapes=small",
+    ],
+    tags = [],
+    target_backends_and_drivers = [
+        ("vulkan-spirv", "vulkan"),
+    ],
+    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
+    test_type = "matmul",
+) for (lhs_rhs_type, acc_type) in (
+    [
+        ("f32", "f32"),
+    ]
+)]
+
 ###########################################################################
 ##
 ## VMVX backend

diff --git a/tests/e2e/matmul/CMakeLists.txt b/tests/e2e/matmul/CMakeLists.txt
@@ -1730,6 +1730,102 @@ iree_generated_e2e_runner_test(
     "arm_64:bf16:+bf16"
 )
 
+iree_generated_e2e_runner_test(
+  NAME
+    e2e_matmul_vmvx_experimental_dt_f32_f32
+  TEST_TYPE
+    matmul
+  GENERATOR
+    "generate_e2e_matmul_tests.py"
+  GENERATOR_ARGS
+    "--lhs_rhs_type=f32"
+    "--acc_type=f32"
+    "--shapes=small"
+  TEST_RUNNER
+    iree_tools_testing_e2e_iree-e2e-matmul-test
+  TARGET_BACKENDS
+    "vmvx"
+  DRIVERS
+    "local-task"
+  COMPILER_FLAGS
+    "--iree-opt-data-tiling"
+    "--iree-global-opt-enable-early-materialization=false"
+  LABELS
+
+)
+
+iree_generated_e2e_runner_test(
+  NAME
+    e2e_matmul_vmvx_experimental_dt_uk_f32_f32
+  TEST_TYPE
+    matmul
+  GENERATOR
+    "generate_e2e_matmul_tests.py"
+  GENERATOR_ARGS
+    "--lhs_rhs_type=f32"
+    "--acc_type=f32"
+    "--shapes=small"
+  TEST_RUNNER
+    iree_tools_testing_e2e_iree-e2e-matmul-test
+  TARGET_BACKENDS
+    "vmvx"
+  DRIVERS
+    "local-task"
+  COMPILER_FLAGS
+    "--iree-opt-data-tiling"
+    "--iree-global-opt-enable-early-materialization=false"
+  LABELS
+
+)
+
+iree_generated_e2e_runner_test(
+  NAME
+    e2e_matmul_cuda_experimental_dt_f32_f32
+  TEST_TYPE
+    matmul
+  GENERATOR
+    "generate_e2e_matmul_tests.py"
+  GENERATOR_ARGS
+    "--lhs_rhs_type=f32"
+    "--acc_type=f32"
+    "--shapes=small"
+  TEST_RUNNER
+    iree_tools_testing_e2e_iree-e2e-matmul-test
+  TARGET_BACKENDS
+    "cuda"
+  DRIVERS
+    "cuda"
+  COMPILER_FLAGS
+    "--iree-opt-data-tiling"
+    "--iree-global-opt-enable-early-materialization=false"
+  LABELS
+
+)
+
+iree_generated_e2e_runner_test(
+  NAME
+    e2e_matmul_spirv_experimental_dt_f32_f32
+  TEST_TYPE
+    matmul
+  GENERATOR
+    "generate_e2e_matmul_tests.py"
+  GENERATOR_ARGS
+    "--lhs_rhs_type=f32"
+    "--acc_type=f32"
+    "--shapes=small"
+  TEST_RUNNER
+    iree_tools_testing_e2e_iree-e2e-matmul-test
+  TARGET_BACKENDS
+    "vulkan-spirv"
+  DRIVERS
+    "vulkan"
+  COMPILER_FLAGS
+    "--iree-opt-data-tiling"
+    "--iree-global-opt-enable-early-materialization=false"
+  LABELS
+
+)
+
 iree_generated_e2e_runner_test(
   NAME
     e2e_matmul_vmvx_dt_uk_i8_small
@@ -2251,6 +2347,35 @@ iree_generated_e2e_runner_test(
     "requires-gpu-cdna3"
 )
 
+iree_generated_e2e_runner_test(
+  NAME
+    e2e_matmul_cdna_experimental_dt_f32_f32
+  TEST_TYPE
+    matmul
+  GENERATOR
+    "generate_e2e_matmul_tests.py"
+  GENERATOR_ARGS
+    "--lhs_rhs_type=f32"
+    "--acc_type=f32"
+    "--shapes=small"
+  TEST_RUNNER
+    iree_tools_testing_e2e_iree-e2e-matmul-test
+  TARGET_BACKENDS
+    "rocm"
+  DRIVERS
+    "hip"
+  COMPILER_FLAGS
+    ${IREE_HIP_TEST_COMPILER_FLAGS}
+    "--iree-opt-data-tiling"
+    "--iree-global-opt-enable-early-materialization=false"
+  LABELS
+    "noasan"
+    "nomsan"
+    "notsan"
+    "noubsan"
+    "requires-gpu-cdna3"
+)
+
 elseif(IREE_HIP_TEST_TARGET_CHIP MATCHES "^gfx11")
 
 unset(IREE_HIP_TEST_COMPILER_FLAGS)
@@ -2320,4 +2445,34 @@ iree_generated_e2e_runner_test(
     "noubsan"
     "requires-gpu-rdna3"
 )
+
+iree_generated_e2e_runner_test(
+  NAME
+    e2e_matmul_rdna3_experimental_dt_f32_f32
+  TEST_TYPE
+    matmul
+  GENERATOR
+    "generate_e2e_matmul_tests.py"
+  GENERATOR_ARGS
+    "--lhs_rhs_type=f32"
+    "--acc_type=f32"
+    "--shapes=small"
+  TEST_RUNNER
+    iree_tools_testing_e2e_iree-e2e-matmul-test
+  TARGET_BACKENDS
+    "rocm"
+  DRIVERS
+    "hip"
+  COMPILER_FLAGS
+    ${IREE_HIP_TEST_COMPILER_FLAGS}
+    "--iree-opt-data-tiling"
+    "--iree-global-opt-enable-early-materialization=false"
+  LABELS
+    "noasan"
+    "nomsan"
+    "notsan"
+    "noubsan"
+    "requires-gpu-rdna3"
+)
+
 endif()