Skip to content

[mlir][amdgpu] Add tensor load store operations #172686

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
amd-eochoalo merged 3 commits into from
Dec 17, 2025
Merged

[mlir][amdgpu] Add tensor load store operations #172686

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
78af3c4
Reapply "[mlir][amdgpu] Add tensor load store operations (#170918)" (…
amd-eochoalo Dec 17, 2025
dcc7b7e
Make type materialization more restrictive
amd-eochoalo Dec 17, 2025
2f55093
comment
amd-eochoalo Dec 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1587,4 +1587,35 @@ def AMDGPU_MakeDmaDescriptorOp : AMDGPU_MakeDescriptorOp<"make_dma_descriptor">

}

def AMDGPU_TensorLoadToLDSOp :
AMDGPU_Op<"tensor_load_to_lds", [MemoryEffects<[MemWrite, MemRead]>]>,
Arguments<(ins AMDGPU_TDMDescriptorType: $desc)> {
let summary = "Load tensors from global memory to LDS.";
let description = [{
Load tensors of up to five dimensions from global memory to LDS.

This operation was introduced in gfx1250.
}];

let assemblyFormat = [{
$desc attr-dict `:` qualified(type($desc))
}];
}

def AMDGPU_TensorStoreFromLDSOp :
AMDGPU_Op<"tensor_store_from_lds", [MemoryEffects<[MemWrite, MemRead]>]>,
Arguments<(ins AMDGPU_TDMDescriptorType: $desc)> {

let summary = "Store tensors from LDS to global memory.";
let description = [{
Store tensors of up to five dimensions from LDS to global memory.

This operation was introduced in gfx1250.
}];

let assemblyFormat = [{
$desc attr-dict `:` qualified(type($desc))
}];
}

#endif // AMDGPU
66 changes: 59 additions & 7 deletions mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3218,11 +3218,6 @@ struct AMDGPULowerDescriptor : public ConvertOpToLLVMPattern<DescriptorOp> {

Location loc = op.getLoc();

IntegerType i32 = rewriter.getI32Type();
[[maybe_unused]] Type v4i32 =
this->typeConverter->convertType(VectorType::get(4, i32));
assert(v4i32 && "expected type conversion to succeed");

SmallVector<Value> consts;
for (int64_t i = 0; i < 8; ++i)
consts.push_back(createI32Constant(rewriter, loc, i));
Expand All @@ -3237,6 +3232,32 @@ struct AMDGPULowerDescriptor : public ConvertOpToLLVMPattern<DescriptorOp> {
}
};

template <typename SourceOp, typename TargetOp>
struct AMDGPUTensorLoadStoreOpLowering
: public ConvertOpToLLVMPattern<SourceOp> {
using ConvertOpToLLVMPattern<SourceOp>::ConvertOpToLLVMPattern;
using Adaptor = typename ConvertOpToLLVMPattern<SourceOp>::OneToNOpAdaptor;
AMDGPUTensorLoadStoreOpLowering(const LLVMTypeConverter &converter,
Chipset chipset)
: ConvertOpToLLVMPattern<SourceOp>(converter), chipset(chipset) {}
Chipset chipset;

LogicalResult
matchAndRewrite(SourceOp op, Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
if (chipset < kGfx1250)
return op->emitOpError("is only supported on gfx1250");

ValueRange desc = adaptor.getDesc();
rewriter.replaceOpWithNewOp<TargetOp>(op, desc[0], desc[1], desc[2],
desc[3], /*cachePolicy=*/0,
/*alias_scopes=*/nullptr,
/*noalias_scopes=*/nullptr,
/*tbaa=*/nullptr);
return success();
}
};

struct ConvertAMDGPUToROCDLPass
: public impl::ConvertAMDGPUToROCDLPassBase<ConvertAMDGPUToROCDLPass> {
using Base::Base;
Expand Down Expand Up @@ -3306,6 +3327,33 @@ void mlir::populateAMDGPUTypeAndAttributeConversions(
Type i32 = IntegerType::get(type.getContext(), 32);
return typeConverter.convertType(VectorType::get(4, i32));
});
typeConverter.addConversion(
[&](TDMDescriptorType type,
SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
Type i32 = IntegerType::get(type.getContext(), 32);
Type v4i32 = typeConverter.convertType(VectorType::get(4, i32));
Type v8i32 = typeConverter.convertType(VectorType::get(8, i32));
llvm::append_values(result, v4i32, v8i32, v4i32, v4i32);
return success();
});

auto addUnrealizedCast = [](OpBuilder &builder, TypeRange types,
ValueRange inputs,
Location loc) -> SmallVector<Value> {
// Only create unrealized_conversion_cast for TDMDescriptorType.
// All other types which are not expected, should be
// materialized by other target materialization functions.
if (inputs.size() != 1)
return {};

if (!isa<TDMDescriptorType>(inputs[0].getType()))
Copy link
Contributor

@Hardcode84 Hardcode84 Dec 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add some comments?

Copy link
Contributor Author

@amd-eochoalo amd-eochoalo Dec 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm writing some comments that will be added in the commit, but maybe it would also help to type it here for some discussion.

The ROCM integration tests broke because I added a materialization cast that uses unrealized_conversion_cast.

  auto addUnrealizedCast = [](OpBuilder &builder, TypeRange types,
                              ValueRange inputs,
                              Location loc) -> SmallVector<Value> {
    auto cast = UnrealizedConversionCastOp::create(builder, loc, types, inputs);
    return cast.getResults();
  };

  typeConverter.addTargetMaterialization(addUnrealizedCast);

This target materialization was intended only for TDMDescriptorType. Without it, we would get errors of the type:

error: unexpected error: failed to legalize unresolved materialization from ('!amdgpu.tdm_descriptor') to ('vector<4xi32>', 'vector<8xi32>', 'vector<4xi32>', 'vector<4xi32>') that remained live after conversion
  amdgpu.tensor_store_from_lds %desc : !amdgpu.tdm_descriptor

But this target materialization is applied too liberally and would be applied to other items like memrefs. (This is taken from the failed tests.)

# | <stdin>:11:12: error: LLVM Translation failed for operation: builtin.unrealized_conversion_cast
# |       %0 = builtin.unrealized_conversion_cast %arg2 : !llvm.ptr to !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
# |            ^
# | <stdin>:11:12: note: see current operation: %0 = "builtin.unrealized_conversion_cast"(%arg2) : (!llvm.ptr) -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>

I originally thought that the need for this target materialization may be a bug in OneToNOpAdaptor since this explicit target materialization was not needed for the regular OpAdaptor. But I asked about this and became convinced that adding the explicit target materialization is ok https://discord.com/channels/636084430946959380/642426447167881246/1445463097262211203

So, in order to avoid adding unrealized_conversion_cast everywhere, we just make it more restrictive and specify its use only in the case of TDMDescriptorType.

I'm still not too sure why it is explicitly needed when it is not needed for TDMBaseType (which does not use OneToNOpAdaptor).

Copy link
Contributor Author

@amd-eochoalo amd-eochoalo Dec 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added comment here : 2f55093 (I'm not too sure how much to add there, the comment above is a little bit more meta than one suitable for an in code comment in my opinion. Let me know what you think! Happy to make edits.)

return {};

auto cast = UnrealizedConversionCastOp::create(builder, loc, types, inputs);
return cast.getResults();
};

typeConverter.addTargetMaterialization(addUnrealizedCast);
}

void mlir::populateAMDGPUToROCDLConversionPatterns(LLVMTypeConverter &converter,
Expand Down Expand Up @@ -3336,7 +3384,11 @@ void mlir::populateAMDGPUToROCDLConversionPatterns(LLVMTypeConverter &converter,
AMDGPUMakeDmaBaseLowering<MakeDmaBaseOp>,
AMDGPUMakeDmaBaseLowering<MakeGatherDmaBaseOp>,
AMDGPULowerDescriptor<MakeDmaDescriptorOp>,
AMDGPULowerDescriptor<MakeGatherDmaDescriptorOp>>(converter,
chipset);
AMDGPULowerDescriptor<MakeGatherDmaDescriptorOp>,
AMDGPUTensorLoadStoreOpLowering<TensorLoadToLDSOp,
ROCDL::TensorLoadToLDSOp>,
AMDGPUTensorLoadStoreOpLowering<TensorStoreFromLDSOp,
ROCDL::TensorStoreFromLDSOp>>(
converter, chipset);
patterns.add<AMDGPUSwizzleBitModeLowering>(converter);
}
18 changes: 18 additions & 0 deletions mlir/test/Conversion/AMDGPUToROCDL/gfx1250.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -773,6 +773,24 @@ func.func @make_dma_descriptor_workgroup_mask(%base: !amdgpu.tdm_base<i32>, %wg_
func.return %descriptor : !amdgpu.tdm_descriptor
}

// CHECK-LABEL: func @tensor_load_to_lds
// CHECK-SAME: (%[[DESC:.+]]: !amdgpu.tdm_descriptor)
func.func @tensor_load_to_lds(%desc: !amdgpu.tdm_descriptor) {
// CHECK: %[[DGROUPS:.+]]:4 = builtin.unrealized_conversion_cast %[[DESC]]
// CHECK: rocdl.tensor.load.to.lds %[[DGROUPS]]#0, %[[DGROUPS]]#1, %[[DGROUPS]]#2, %[[DGROUPS]]#3 cachepolicy 0 : vector<4xi32>, vector<8xi32>
amdgpu.tensor_load_to_lds %desc : !amdgpu.tdm_descriptor
func.return
}

// CHECK-LABEL: func @tensor_store_from_lds
// CHECK-SAME: (%[[DESC:.+]]: !amdgpu.tdm_descriptor)
func.func @tensor_store_from_lds(%desc: !amdgpu.tdm_descriptor) {
// CHECK: %[[DGROUPS:.+]]:4 = builtin.unrealized_conversion_cast %[[DESC]]
// CHECK: rocdl.tensor.store.from.lds %[[DGROUPS]]#0, %[[DGROUPS]]#1, %[[DGROUPS]]#2, %[[DGROUPS]]#3 cachepolicy 0 : vector<4xi32>, vector<8xi32>
amdgpu.tensor_store_from_lds %desc : !amdgpu.tdm_descriptor
func.return
}

// -----

// CHECK-LABEL: func @make_gather_dma_descriptor
Expand Down