[Mlir-commits] [mlir] [mlir][amdgpu] Add tensor load store operations (PR #170918)

[Krzysztof Drewniak]

================
@@ -440,6 +440,126 @@ func.func @make_dma_descriptor_workgroup_mask(%base: !amdgpu.tdm_base<i32>, %wg_
   // CHECK: %[[DGROUP1:.+]] = llvm.insertelement %[[SGPR7]], %[[DGROUP1_6]][%[[C7]] : i32]
 
   // CHECK: %[[DGROUPS:.+]] = builtin.unrealized_conversion_cast %[[DGROUP0]], %[[DGROUP1]] : vector<4xi32>, vector<8xi32> to !amdgpu.tdm_descriptor
-  %descriptor = amdgpu.make_dma_descriptor %base globalSize [128, 64] globalStride [64, 1] sharedSize [128, 64] workgroupMask %wg_mask earlyTimeout %timeout : !amdgpu.tdm_base<i32> -> !amdgpu.tdm_descriptor
-  func.return %descriptor : !amdgpu.tdm_descriptor
+  %descriptor = amdgpu.make_dma_descriptor %base globalSize [128, 64] globalStride [64, 1] sharedSize [128, 64] workgroupMask %wg_mask earlyTimeout %timeout : !amdgpu.tdm_base<i32> -> !amdgpu.tdm_descriptor<2>
+  func.return %descriptor : !amdgpu.tdm_descriptor<2>
+}
+
+// CHECK-LABEL: func @tensor_load_to_lds_d2
+// CHECK-SAME: (%[[DESC:.+]]: !amdgpu.tdm_descriptor<2>)
+func.func @tensor_load_to_lds_d2(%desc: !amdgpu.tdm_descriptor<2>) {
+  // CHECK: %[[DGROUPS:.+]]:2 = builtin.unrealized_conversion_cast %[[DESC]]
+  // CHECK: rocdl.tensor.load.to.lds.d2 %[[DGROUPS]]#0, %[[DGROUPS]]#1 cachepolicy 0 : vector<4xi32>, vector<8xi32>
+  amdgpu.tensor_load_to_lds %desc : !amdgpu.tdm_descriptor<2>
+
+  // CHECK: rocdl.tensor.load.to.lds.d2 %[[DGROUPS]]#0, %[[DGROUPS]]#1 cachepolicy 0 : vector<4xi32>, vector<8xi32>
+  amdgpu.tensor_load_to_lds %desc { cache_scope = #amdgpu.cache_scope<workgroup> } : !amdgpu.tdm_descriptor<2>
----------------
krzysz00 wrote:

So, given that we've been cleared for the compiler rewrite that just lets us replace zeros with SGPR NULL, I'm actually going to vote against putting a size parameter on here

https://github.com/llvm/llvm-project/pull/170918