Amend from conversion to dialect rewrite pattern

Author: jerryyin

mlir/include/mlir/Conversion/Passes.h

@@ -73,7 +73,6 @@
 #include "mlir/Conversion/TosaToTensor/TosaToTensor.h"
 #include "mlir/Conversion/UBToLLVM/UBToLLVM.h"
 #include "mlir/Conversion/UBToSPIRV/UBToSPIRV.h"
-#include "mlir/Conversion/VectorToAMDGPU/VectorToAMDGPU.h"
 #include "mlir/Conversion/VectorToArmSME/VectorToArmSME.h"
 #include "mlir/Conversion/VectorToGPU/VectorToGPU.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.h"

mlir/include/mlir/Conversion/Passes.td

@@ -1333,16 +1333,6 @@ def ConvertVectorToArmSMEPass : Pass<"convert-vector-to-arm-sme"> {
   let dependentDialects = ["arm_sme::ArmSMEDialect", "arm_sve::ArmSVEDialect"];
 }
 
-//===----------------------------------------------------------------------===//
-// VectorToAMDGPU
-//===----------------------------------------------------------------------===//
-
-def ConvertVectorToAMDGPUPass : Pass<"convert-vector-to-amdgpu"> {
-  let summary = "Lower the operations from the vector dialect into the AMDGPU "
-                "dialect";
-  let dependentDialects = ["vector::VectorDialect"];
-} 
-
 //===----------------------------------------------------------------------===//
 // ArmSMEToSCF
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Conversion/VectorToAMDGPU/VectorToAMDGPU.h

@@ -22,6 +22,7 @@ namespace amdgpu {
 
 #define GEN_PASS_DECL_AMDGPUEMULATEATOMICSPASS
 #define GEN_PASS_DECL_AMDGPURESOLVESTRIDEDMETADATAPASS
+#define GEN_PASS_DECL_AMDGPUTRANSFERREADTOLOADPASS
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/AMDGPU/Transforms/Passes.h.inc"
 
@@ -30,6 +31,9 @@ void populateAmdgpuEmulateAtomicsPatterns(ConversionTarget &target,
                                           Chipset chipset);
 
 void populateAmdgpuResolveStridedMetadataPatterns(RewritePatternSet &patterns);
+
+void populateAmdgpuTransferReadToLoadPatterns(RewritePatternSet &patterns);
+
 } // namespace amdgpu
 } // namespace mlir
 

mlir/include/mlir/Dialect/AMDGPU/Transforms/Passes.h

@@ -51,4 +51,18 @@ def AmdgpuResolveStridedMetadataPass : Pass<"amdgpu-resolve-strided-metadata"> {
   ];
 }
 
+def AmdgpuTransferReadToLoadPass : Pass<"amdgpu-transfer-read-to-load"> {
+  let summary = "Lower the operations from the vector transfer_read to vector load";
+  let description = [{
+    This pass creates a transfer read op lowering. A vector trasfer read op
+    will be lowered to a combination of vector.load, arith.select and
+    vector.broadcast.
+
+    This pattern will make it possible for masked transfer_read to be lowered
+    towards buffer load with bounds check, allowing a more optimized global
+    load accessing pattern compared with existing implementation of
+    llvm.intr.masked.load on vectors.
+  }];
+  let dependentDialects = [];
+}
 #endif // MLIR_DIALECT_AMDGPU_TRANSFORMS_PASSES_TD_

mlir/include/mlir/Dialect/AMDGPU/Transforms/Passes.td

@@ -66,7 +66,6 @@ add_subdirectory(TosaToSCF)
 add_subdirectory(TosaToTensor)
 add_subdirectory(UBToLLVM)
 add_subdirectory(UBToSPIRV)
-add_subdirectory(VectorToAMDGPU)
 add_subdirectory(VectorToArmSME)
 add_subdirectory(VectorToGPU)
 add_subdirectory(VectorToLLVM)

mlir/lib/Conversion/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_mlir_dialect_library(MLIRAMDGPUTransforms
   EmulateAtomics.cpp
   ResolveStridedMetadata.cpp
+  TransferReadToLoad.cpp
 
   ADDITIONAL_HEADER_DIRS
   {$MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/AMDGPU/Transforms

mlir/lib/Conversion/VectorToAMDGPU/CMakeLists.txt

@@ -1,12 +1,12 @@
-//===- VectorToAMDGPU.cpp - Vector to AMDGPU dialect conversion ---------===//
+//===- TransferReadToLoad.cpp - Lowers masked transfer read to load -------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Conversion/VectorToAMDGPU/VectorToAMDGPU.h"
+#include "mlir/Dialect/AMDGPU/Transforms/Passes.h"
 
 #include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
@@ -17,12 +17,13 @@
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/WalkPatternRewriteDriver.h"
 
-namespace mlir {
-#define GEN_PASS_DEF_CONVERTVECTORTOAMDGPUPASS
-#include "mlir/Conversion/Passes.h.inc"
-} // namespace mlir
+namespace mlir::amdgpu {
+#define GEN_PASS_DEF_AMDGPUTRANSFERREADTOLOADPASS
+#include "mlir/Dialect/AMDGPU/Transforms/Passes.h.inc"
+} // namespace mlir::amdgpu
 
 using namespace mlir;
+using namespace mlir::amdgpu;
 
 /// This pattern supports lowering of:
 /// `vector.transfer_read` to a combination of `vector.load`, `arith.select` and
@@ -55,8 +56,11 @@ static LogicalResult transferPreconditions(
     return rewriter.notifyMatchFailure(xferOp, "not a memref source");
 
   Attribute addrSpace = memRefType.getMemorySpace();
-  if (!addrSpace || dyn_cast<amdgpu::AddressSpaceAttr>(addrSpace).getValue() !=
-                        amdgpu::AddressSpace::FatRawBuffer)
+  if (!addrSpace || !dyn_cast<amdgpu::AddressSpaceAttr>(addrSpace))
+    return rewriter.notifyMatchFailure(xferOp, "no address space");
+
+  if (dyn_cast<amdgpu::AddressSpaceAttr>(addrSpace).getValue() !=
+      amdgpu::AddressSpace::FatRawBuffer)
     return rewriter.notifyMatchFailure(xferOp, "not in buffer address space");
 
   // Non-unit strides are handled by VectorToSCF.
@@ -134,16 +138,17 @@ struct TransferReadLowering final : OpRewritePattern<vector::TransferReadOp> {
 
 } // namespace
 
-void mlir::populateVectorToAMDGPUConversionPatterns(
+void mlir::amdgpu::populateAmdgpuTransferReadToLoadPatterns(
     RewritePatternSet &patterns) {
   patterns.add<TransferReadLowering>(patterns.getContext());
 }
 
-struct ConvertVectorToAMDGPUPass final
-    : impl::ConvertVectorToAMDGPUPassBase<ConvertVectorToAMDGPUPass> {
+struct AmdgpuTransferReadToLoadPass final
+    : amdgpu::impl::AmdgpuTransferReadToLoadPassBase<
+          AmdgpuTransferReadToLoadPass> {
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
-    populateVectorToAMDGPUConversionPatterns(patterns);
+    populateAmdgpuTransferReadToLoadPatterns(patterns);
     walkAndApplyPatterns(getOperation(), std::move(patterns));
   }
 };

mlir/lib/Dialect/AMDGPU/Transforms/CMakeLists.txt

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s --convert-vector-to-amdgpu --split-input-file | FileCheck %s
+// RUN: mlir-opt %s --amdgpu-transfer-read-to-load --split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @transfer_to_maskedload_fatrawbuffer(
 // CHECK-SAME: %[[ARG0:.*]]: memref<8x8xf32, #amdgpu.address_space<fat_raw_buffer>>
@@ -32,6 +32,21 @@ func.func @transfer_to_maskedload_regular(%mem : memref<8x8xf32>, %idx : index,
 
 // -----
 
+// CHECK-LABEL: func @transfer_to_maskedload_addrspace(
+// CHECK-SAME: %[[ARG0:.*]]: memref<8x8xf32, #gpu.address_space<workgroup>>
+// CHECK-SAME: %[[ARG1:.*]]: index
+// CHECK-SAME: %[[ARG2:.*]]: vector<4xi1>
+func.func @transfer_to_maskedload_addrspace(%mem : memref<8x8xf32, #gpu.address_space<workgroup>>, %idx : index, %mask : vector<4xi1>) -> vector<4xf32> {
+  %cf0 = arith.constant 0.0 : f32
+  %res = vector.transfer_read %mem[%idx, %idx], %cf0, %mask {in_bounds = [true]} : memref<8x8xf32, #gpu.address_space<workgroup>>, vector<4xf32>
+  return %res : vector<4xf32>
+}
+// CHECK: %[[CST:.*]] = arith.constant 0.0
+// CHECK: %[[RES:.*]] = vector.transfer_read %arg0[%arg1, %arg1], %[[CST]], %arg2 {in_bounds = [true]} : memref<8x8xf32, #gpu.address_space<workgroup>>, vector<4xf32>
+// CHECK: return %[[RES]] : vector<4xf32>
+
+// -----
+
 // CHECK-LABEL: func @transfer_broadcasting(
 // CHECK-SAME: %[[ARG0:.*]]: memref<8x8xf32, #amdgpu.address_space<fat_raw_buffer>>
 // CHECK-SAME: %[[ARG1:.*]]: index