[mlir][AMDGPU] Explicitly truncate memory addresses in buffer ops

[krzysz00]

Summary

As a percaution, truncate memory addresses passed to kernels to 48 bits,
since bits 48-63 of the buffer descriptor are used for the stride field
and, on gfx10, to control swizzling.

[ThomasRaoux]

Could you add a test?

[krzysz00]

Added test

[ThomasRaoux]

LGTM

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

matchAndRewrite(GpuOp gpuOp, typename GpuOp::Adaptor adaptor,
     // bit 63: enable swizzles (always off for raw buffers)
     // bits 64-95 (word 2): Number of records, units of stride
     // bits 96-127 (word 3): See below
 
     Type llvm4xI32 = this->typeConverter->convertType(VectorType::get(4, i32));
     MemRefDescriptor memrefDescriptor(memref);
     Type llvmI64 = this->typeConverter->convertType(rewriter.getI64Type());
     Type llvm2xI32 = this->typeConverter->convertType(VectorType::get(2, i32));
+    Value c32I64 = rewriter.create<LLVM::ConstantOp>(
+        loc, llvmI64, rewriter.getI64IntegerAttr(32));
 
     Value resource = rewriter.create<LLVM::UndefOp>(loc, llvm4xI32);
 
     Value ptr = memrefDescriptor.alignedPtr(rewriter, loc);
     Value ptrAsInt = rewriter.create<LLVM::PtrToIntOp>(loc, llvmI64, ptr);
-    Value ptrAsInts =
-        rewriter.create<LLVM::BitcastOp>(loc, llvm2xI32, ptrAsInt);
-    for (int64_t i = 0; i < 2; ++i) {
-      Value idxConst = this->createIndexConstant(rewriter, loc, i);
-      Value part =
-          rewriter.create<LLVM::ExtractElementOp>(loc, ptrAsInts, idxConst);
-      resource = rewriter.create<LLVM::InsertElementOp>(
-          loc, llvm4xI32, resource, part, idxConst);
-    }
+    Value lowHalf = rewriter.create<LLVM::TruncOp>(loc, llvmI32, ptrAsInt);
+        loc, llvm4xI32, resource, lowHalf,
+        this->createIndexConstant(rewriter, loc, 0));
+
+    // Bits 48-63 are used both for the stride of the buffer and (on gfx10) for
+    // enabling swizzling. Prevent the high bits of pointers from accidentally
+    // setting those flags.
+    Value highHalfShifted = rewriter.create<LLVM::TruncOp>(
+        loc, llvmI32, rewriter.create<LLVM::LShrOp>(loc, ptrAsInt, c32I64));
+    Value highHalfTruncated = rewriter.create<LLVM::AndOp>(
+        loc, llvmI32, highHalfShifted,
+        createI32Constant(rewriter, loc, 0x0000ffff));
+    resource = rewriter.create<LLVM::InsertElementOp>(
+        loc, llvm4xI32, resource, highHalfTruncated,
+        this->createIndexConstant(rewriter, loc, 1));
 
     Value numRecords;
     if (memrefType.hasStaticShape()) {
       numRecords = createI32Constant(
           rewriter, loc,
           static_cast<int32_t>(memrefType.getNumElements() * elementByteWidth));
     } else {
       Value maxIndex;

mlir/test/Conversion/AMDGPUToROCDL/amdgpu-to-rocdl.mlir

 // RUN: mlir-opt %s -convert-amdgpu-to-rocdl=chipset=gfx908 | FileCheck %s
 // RUN: mlir-opt %s -convert-amdgpu-to-rocdl=chipset=gfx1030 | FileCheck %s --check-prefix=RDNA
 
 // CHECK-LABEL: func @gpu_gcn_raw_buffer_load_i32
 func.func @gpu_gcn_raw_buffer_load_i32(%buf: memref<64xi32>, %idx: i32) -> i32 {
+  // CHECK: %[[ptr:.*]] = llvm.ptrtoint
+  // CHECK: %[[lowHalf:.*]] = llvm.trunc %[[ptr]] : i64 to i32
+  // CHECK: %[[resource_1:.*]] = llvm.insertelement %[[lowHalf]]
+  // CHECK: %[[highHalfI64:.*]] = llvm.lshr %[[ptr]]
+  // CHECK: %[[highHalfI32:.*]] = llvm.trunc %[[highHalfI64]] : i64 to i32
+  // CHECK: %[[highHalf:.*]] = llvm.and %[[highHalfI32]], %{{.*}} : i32
+  // CHECK: %[[resource_2:.*]] = llvm.insertelement %[[highHalf]], %[[resource_1]]
   // CHECK: %[[numRecords:.*]] = llvm.mlir.constant(256 : i32)
-  // CHECK: llvm.insertelement{{.*}}%[[numRecords]]
+  // CHECK: %[[resource_3:.*]] = llvm.insertelement %[[numRecords]], %[[resource_2]]
   // CHECK: %[[word3:.*]] = llvm.mlir.constant(159744 : i32)
   // RDNA: %[[word3:.*]] = llvm.mlir.constant(822243328 : i32)
-  // CHECK: %[[resource:.*]] = llvm.insertelement{{.*}}%[[word3]]
+  // CHECK: %[[resource:.*]] = llvm.insertelement %[[word3]], %[[resource_3]]
   // CHECK: %[[ret:.*]] = rocdl.raw.buffer.load %[[resource]], %{{.*}}, %{{.*}}, %{{.*}} : i32
   // CHECK: return %[[ret]]
   %0 = amdgpu.raw_buffer_load {boundsCheck = true} %buf[%idx] : memref<64xi32>, i32 -> i32
   func.return %0 : i32
 }
 
 // CHECK-LABEL: func @gpu_gcn_raw_buffer_load_i32_oob_off
 func.func @gpu_gcn_raw_buffer_load_i32_oob_off(%buf: memref<64xi32>, %idx: i32) -> i32 {