[mlir][amdgpu] Add tensor load store operations

[amd-eochoalo]

Reland #170918

This PR differs from the original one by making the target materialization more restrictive.

[llvmbot]

@llvm/pr-subscribers-mlir-amdgpu
@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-backend-amdgpu

@llvm/pr-subscribers-mlir-gpu

Author: Erick Ochoa Lopez (amd-eochoalo)

Changes

Reland #170918

This PR differs from the original one by making the target materialization more restrictive.

---

Full diff: https://github.com/llvm/llvm-project/pull/172686.diff

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td (+31)
• (modified) mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp (+56-7)
• (modified) mlir/test/Conversion/AMDGPUToROCDL/gfx1250.mlir (+18)

diff --git a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td
index 96f5f5c6f1a3f..4865dc13f324b 100644
--- a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td
+++ b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td
@@ -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
diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
index 541bb02d79eae..706b17f1d2b45 100644
--- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
+++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
@@ -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));
@@ -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;
@@ -3306,6 +3327,30 @@ 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> {
+    if (inputs.size() != 1)
+      return {};
+
+    if (!isa<TDMDescriptorType>(inputs[0].getType()))
+      return {};
+
+    auto cast = UnrealizedConversionCastOp::create(builder, loc, types, inputs);
+    return cast.getResults();
+  };
+
+  typeConverter.addTargetMaterialization(addUnrealizedCast);
 }
 
 void mlir::populateAMDGPUToROCDLConversionPatterns(LLVMTypeConverter &converter,
@@ -3336,7 +3381,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);
 }
diff --git a/mlir/test/Conversion/AMDGPUToROCDL/gfx1250.mlir b/mlir/test/Conversion/AMDGPUToROCDL/gfx1250.mlir
index 4979e85785970..e62db9ff571bf 100644
--- a/mlir/test/Conversion/AMDGPUToROCDL/gfx1250.mlir
+++ b/mlir/test/Conversion/AMDGPUToROCDL/gfx1250.mlir
@@ -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

[Hardcode84]

Can you add some comments?

[amd-eochoalo]

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).

[amd-eochoalo]

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.)

[krzysz00]

No objection here either