From b740597021919b5ffe626df014f872730c6dc50e Mon Sep 17 00:00:00 2001 From: Stanley Winata Date: Mon, 21 Jul 2025 15:49:06 -0700 Subject: [PATCH 1/2] [MLIR][AMDGPU] Use Attr for resetOffset + boundsCheck in RawBufferCastOp In order to access and modify resetOffset and boundsCheck of RawBufferCastOp in pythonic binding, we will have to use Attrs instead of Property. This is because we do not have python binding support for property yet. We should move back to property once we add pythonic binding support for it. Signed-off-by: Stanley Winata --- mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td | 4 +-- mlir/test/python/dialects/amdgpu.py | 27 ++++++++++++++++++- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td index 80959ffbaf426..e2ca1ba9ae5e6 100644 --- a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td +++ b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td @@ -237,8 +237,8 @@ def AMDGPU_FatRawBufferCastOp : Arguments<(ins AnyMemRef:$source, Optional:$validBytes, Optional>:$cacheSwizzleStride, - DefaultValuedProp:$boundsCheck, - UnitProp:$resetOffset)>, + DefaultValuedAttr:$boundsCheck, + UnitAttr:$resetOffset)>, Results<(outs AnyMemRef:$result)> { let summary = "Create a raw buffer fat pointer that matches `memref`"; let description = [{ diff --git a/mlir/test/python/dialects/amdgpu.py b/mlir/test/python/dialects/amdgpu.py index c8039d494cf81..090888a94d667 100644 --- a/mlir/test/python/dialects/amdgpu.py +++ b/mlir/test/python/dialects/amdgpu.py @@ -2,7 +2,7 @@ # This is just a smoke test that the dialect is functional. from mlir.ir import * -from mlir.dialects import amdgpu, arith, memref +from mlir.dialects import amdgpu, func def constructAndPrintInModule(f): @@ -20,3 +20,28 @@ def constructAndPrintInModule(f): def testSmoke(): # CHECK: amdgpu.lds_barrier amdgpu.LDSBarrierOp() + + +# CHECK-LABEL: testFatRawBufferCastOpParams +@constructAndPrintInModule +def testFatRawBufferCastOpParams(): + memref_type = MemRefType.get( + [ShapedType.get_dynamic_size(), ShapedType.get_dynamic_size()], + F32Type.get(), + ) + f = func.FuncOp( + "test_raw_buffer_cast_params", ([memref_type], []) + ) + with InsertionPoint(f.add_entry_block()): + block_args = f.arguments + amdgpu.FatRawBufferCastOp(block_args[0]) + amdgpu.FatRawBufferCastOp(block_args[0], resetOffset=True) + amdgpu.FatRawBufferCastOp(block_args[0], boundsCheck=False) + amdgpu.FatRawBufferCastOp(block_args[0], boundsCheck=False, resetOffset=True) + func.ReturnOp([]) + + #CHECK: func.func @test_raw_buffer_cast_params(%[[ARG0:.+]]: memref) { + #CHECK: amdgpu.fat_raw_buffer_cast %[[ARG0]] : memref to memref> + #CHECK-NEXT: amdgpu.fat_raw_buffer_cast %[[ARG0]] resetOffset : memref to memref> + #CHECK-NEXT: amdgpu.fat_raw_buffer_cast %[[ARG0]] boundsCheck(false) : memref to memref> + #CHECK-NEXT: amdgpu.fat_raw_buffer_cast %[[ARG0]] boundsCheck(false) resetOffset : memref to memref> From f96ff80ed52706e3dc25c2d399f33b609e62dd84 Mon Sep 17 00:00:00 2001 From: Stanley Winata Date: Mon, 21 Jul 2025 16:10:30 -0700 Subject: [PATCH 2/2] add TODO and lint python test Signed-off-by: Stanley Winata --- mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td | 2 ++ mlir/test/python/dialects/amdgpu.py | 14 ++++++-------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td index e2ca1ba9ae5e6..5a53b15a9c679 100644 --- a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td +++ b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td @@ -240,6 +240,8 @@ def AMDGPU_FatRawBufferCastOp : DefaultValuedAttr:$boundsCheck, UnitAttr:$resetOffset)>, Results<(outs AnyMemRef:$result)> { + // TODO: Set `resetOffset` and `boundsCheck` to use `Property` once + // we implemented pythonic binding for `Property`. let summary = "Create a raw buffer fat pointer that matches `memref`"; let description = [{ Wraps the memory pointed to by `source` as a raw buffer fat pointer, or, diff --git a/mlir/test/python/dialects/amdgpu.py b/mlir/test/python/dialects/amdgpu.py index 090888a94d667..b479576dac093 100644 --- a/mlir/test/python/dialects/amdgpu.py +++ b/mlir/test/python/dialects/amdgpu.py @@ -29,9 +29,7 @@ def testFatRawBufferCastOpParams(): [ShapedType.get_dynamic_size(), ShapedType.get_dynamic_size()], F32Type.get(), ) - f = func.FuncOp( - "test_raw_buffer_cast_params", ([memref_type], []) - ) + f = func.FuncOp("test_raw_buffer_cast_params", ([memref_type], [])) with InsertionPoint(f.add_entry_block()): block_args = f.arguments amdgpu.FatRawBufferCastOp(block_args[0]) @@ -40,8 +38,8 @@ def testFatRawBufferCastOpParams(): amdgpu.FatRawBufferCastOp(block_args[0], boundsCheck=False, resetOffset=True) func.ReturnOp([]) - #CHECK: func.func @test_raw_buffer_cast_params(%[[ARG0:.+]]: memref) { - #CHECK: amdgpu.fat_raw_buffer_cast %[[ARG0]] : memref to memref> - #CHECK-NEXT: amdgpu.fat_raw_buffer_cast %[[ARG0]] resetOffset : memref to memref> - #CHECK-NEXT: amdgpu.fat_raw_buffer_cast %[[ARG0]] boundsCheck(false) : memref to memref> - #CHECK-NEXT: amdgpu.fat_raw_buffer_cast %[[ARG0]] boundsCheck(false) resetOffset : memref to memref> + # CHECK: func.func @test_raw_buffer_cast_params(%[[ARG0:.+]]: memref) { + # CHECK: amdgpu.fat_raw_buffer_cast %[[ARG0]] : memref to memref> + # CHECK-NEXT: amdgpu.fat_raw_buffer_cast %[[ARG0]] resetOffset : memref to memref> + # CHECK-NEXT: amdgpu.fat_raw_buffer_cast %[[ARG0]] boundsCheck(false) : memref to memref> + # CHECK-NEXT: amdgpu.fat_raw_buffer_cast %[[ARG0]] boundsCheck(false) resetOffset : memref to memref>