[X86] Transform amx pointer.

[LuoYuanke]

Summary

Front-end bitcast <256 x i32> to x86_amx and generate load/store <256 x i32>*.
In instruction combine pass it transform load/store <256 x i32>* to
load/store x86_amx*. In the amx type lowering pass , we lower the load/store
instructions to amx load/store intrinsics with the stride value 64.

[LuoYuanke]

Rebase.

[LuoYuanke]

Support PHI operation.

[RKSimon]

a few style comments

llvm/lib/Target/X86/X86LowerAMXType.cpp

94

isa<>

122

isa<>

289

'''
if (auto *ST = dyn_cast<StoreInst>(I))
'''

[lebedev.ri]

What's the status here?

[lebedev.ri]

This review seems to be stuck/dead, consider abandoning if no longer relevant.

[LuoYuanke]

We have amx cast instrinsics for vector and x86_amx type and front-end would use amx cast instrinsics instead of bitcast. I don't think we need this patch now.

llvm/lib/Target/X86/X86LowerAMXType.cpp

case Intrinsic::x86_tdpbssd_internal: {
     }
     break;
   }
   }
 
   return std::make_pair(Row, Col);
 }
 
+// %1 = load x86_amx, x86_amx* %0, align 64
+// %2 = call x86_amx @llvm.x86.tdpbssd.internal(%1, %1, %1, ...)
+// -->
+// %1 = call x86_amx @llvm.x86.tileloadd64.internal()
+// %2 = call x86_amx @llvm.x86.tdpbssd.internal(%1, %1, %1, ...)
+static void transformTileLoad(LoadInst *LD) {
+  Value *Row = nullptr, *Col = nullptr;
+  Use &U = *(LD->use_begin());
+  unsigned OpNo = U.getOperandNo();
+  auto *II = cast<IntrinsicInst>(U.getUser());
+  std::tie(Row, Col) = getShape(II, OpNo);
+  IRBuilder<> Builder(LD);
+  // Use the maximun column as stride.
+  Value *Stride = Builder.getInt64(64);
+  Value *I8Ptr =
+      Builder.CreateBitCast(LD->getOperand(0), Builder.getInt8PtrTy());
+  std::array<Value *, 4> Args = {Row, Col, I8Ptr, Stride};
+
+  Value *NewInst =
+      Builder.CreateIntrinsic(Intrinsic::x86_tileloadd64_internal, None, Args);
+  LD->replaceAllUsesWith(NewInst);
+}
+
 // %src = load <256 x i32>, <256 x i32>* %addr, align 64
 // %2 = bitcast <256 x i32> %src to x86_amx
 // -->
 // %2 = call x86_amx @llvm.x86.tileloadd64.internal(i16 %row, i16 %col,
 // i8* %addr, i64 %stride64)
 static void combineLoadBitcast(LoadInst *LD, BitCastInst *Bitcast) {
   Value *Row = nullptr, *Col = nullptr;
   Use &U = *(Bitcast->use_begin());
   unsigned OpNo = U.getOperandNo();
   auto *II = cast<IntrinsicInst>(U.getUser());
   std::tie(Row, Col) = getShape(II, OpNo);
   IRBuilder<> Builder(Bitcast);
   // Use the maximun column as stride.
   Value *Stride = Builder.getInt64(64);
   Value *I8Ptr =
       Builder.CreateBitCast(LD->getOperand(0), Builder.getInt8PtrTy());
   std::array<Value *, 4> Args = {Row, Col, I8Ptr, Stride};
 
if (Bitcast->getType()->isX86_AMXTy()) {
     Builder.CreateIntrinsic(Intrinsic::x86_tilestored64_internal, None, Args);
     Value *NewInst = Builder.CreateLoad(Bitcast->getType(), AllocaAddr);
     Bitcast->replaceAllUsesWith(NewInst);
   }
 
   return true;
 }
 
+// %addr = bitcast <256 x i32>* %tile to x86_amx*
+// store x86_amx %9, x86_amx* %addr, align 64
+// -->
+// call void @llvm.x86.tilestored64.internal(%row, %col, %addr,
+// %stride64, %9)
+static void transformTileStore(StoreInst *ST) {
+  auto *II = cast<IntrinsicInst>(ST->getValueOperand());
+  Value *Row = II->getOperand(0);
+  Value *Col = II->getOperand(1);
+  IRBuilder<> Builder(ST);
+  // Use the maximum column as stride. It must be the same with load
+  // stride.
+  Value *Stride = Builder.getInt64(64);
+  Value *I8Ptr =
+      Builder.CreateBitCast(ST->getOperand(1), Builder.getInt8PtrTy());
+  std::array<Value *, 5> Args = {Row, Col, I8Ptr, Stride,
+                                 ST->getValueOperand()};
+  Builder.CreateIntrinsic(Intrinsic::x86_tilestored64_internal, None, Args);
+}
+
 namespace {
 class X86LowerAMXType {
   Function &Func;
 
 public:
   X86LowerAMXType(Function &F) : Func(F) {}
   bool visit();
 };
 
 bool X86LowerAMXType::visit() {
   SmallVector<Instruction *, 8> DeadInsts;
 
   for (BasicBlock *BB : post_order(&Func)) {
     for (BasicBlock::reverse_iterator II = BB->rbegin(), IE = BB->rend();
          II != IE;) {
       Instruction &Inst = *II++;
       auto *Bitcast = dyn_cast<BitCastInst>(&Inst);
       if (!Bitcast)
         continue;
 
       Value *Src = Bitcast->getOperand(0);
+      Type *Ty = Bitcast->getType();
+
+      if (Ty->isPointerTy() &&
+          cast<PointerType>(Ty)->getElementType()->isX86_AMXTy()) {
+        for (auto UI = Bitcast->use_begin(), UE = Bitcast->use_end();
+             UI != UE;) {
+          Value *I = (UI++)->getUser();
+          auto *LD = dyn_cast<LoadInst>(I);
+          // %0 = bitcast <256 x i32>* %tile to x86_amx*
+          // %1 = load x86_amx, x86_amx* %0, align 64
+          if (LD) {
+            transformTileLoad(LD);
+            DeadInsts.push_back(LD);
+          }
+          auto *ST = dyn_cast<StoreInst>(I);
+          if (ST) {
+            // %addr = bitcast <256 x i32>* %tile to x86_amx*
+            // store x86_amx %9, x86_amx* %addr, align 64
+            // -->
+            // call void @llvm.x86.tilestored64.internal(%row, %col, %addr,
+            // %stride64, %9)
+            transformTileStore(ST);
+            DeadInsts.push_back(ST);
+          }
+        }
+        // If the dst type is <256 x i32>*, it is valid intruction.
+        // %0 = bitcast x86_amx* %tile to <256 x i32>*
+        // %1 = load <256 x i32>, <256 x i32>* %0, align 64
+        // store <256 x i32> %2, <256 x i32>* %0, align 64
+      }
       if (Bitcast->getType()->isX86_AMXTy()) {
         if (Bitcast->user_empty()) {
           DeadInsts.push_back(Bitcast);
           continue;
         }
         LoadInst *LD = dyn_cast<LoadInst>(Src);
         if (!LD) {
           if (transformBitcast(Bitcast))

llvm/test/CodeGen/X86/AMX/amx-type.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -lower-amx-type %s -S | FileCheck %s
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 %struct.__tile_str = type { i16, i16, <256 x i32> }
 
 @buf = dso_local global [1024 x i8] zeroinitializer, align 16
 @buf2 = dso_local global [1024 x i8] zeroinitializer, align 16
 
+define dso_local void @test_amx_store(<256 x i32>* %in, i16 %m, i16 %n, i8 *%buf, i64 %s) #2 {
+; CHECK-LABEL: @test_amx_store(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[T0:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[M:%.*]], i16 [[N:%.*]], i8* [[BUF:%.*]], i64 [[S:%.*]]) [[ATTR3:#.*]]
+; CHECK-NEXT:    [[ADDR:%.*]] = bitcast <256 x i32>* [[IN:%.*]] to x86_amx*
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast x86_amx* [[ADDR]] to i8*
+; CHECK-NEXT:    call void @llvm.x86.tilestored64.internal(i16 [[M]], i16 [[N]], i8* [[TMP0]], i64 64, x86_amx [[T0]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %t0 = call x86_amx @llvm.x86.tileloadd64.internal(i16 %m, i16 %n, i8* %buf, i64 %s) #3
+  %addr = bitcast <256 x i32>* %in to x86_amx*
+  store x86_amx %t0, x86_amx* %addr, align 64
+  ret void
+}
+
+define dso_local void @test_amx_load(<256 x i32>* %in, i16 %m, i16 %n, i8 *%buf, i64 %s) #2 {
+; CHECK-LABEL: @test_amx_load(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[T0:%.*]] = bitcast <256 x i32>* [[IN:%.*]] to x86_amx*
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast x86_amx* [[T0]] to i8*
+; CHECK-NEXT:    [[TMP1:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[M:%.*]], i16 [[N:%.*]], i8* [[TMP0]], i64 64)
+; CHECK-NEXT:    call void @llvm.x86.tilestored64.internal(i16 [[M]], i16 [[N]], i8* [[BUF:%.*]], i64 [[S:%.*]], x86_amx [[TMP1]]) [[ATTR3]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %t0 = bitcast <256 x i32>* %in to x86_amx*
+  %t1 = load x86_amx, x86_amx* %t0, align 64
+  call void @llvm.x86.tilestored64.internal(i16 %m, i16 %n, i8* %buf, i64 %s, x86_amx %t1) #3
+  ret void
+}
+
 ; test bitcast x86_amx to <256 x i32>
 define dso_local void @test_user_empty(i16 %m, i16 %n, i8 *%buf, i64 %s) #2 {
 ; CHECK-LABEL: @test_user_empty(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %t1 = call x86_amx @llvm.x86.tileloadd64.internal(i16 %m, i16 %n, i8* %buf, i64 %s) #3

 }
 
 define dso_local <256 x i32> @test_amx_load_bitcast(<256 x i32>* %in, i16 %m, i16 %n, i8 *%buf, i64 %s) #2 {
 ; CHECK-LABEL: @test_amx_load_bitcast(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[T1:%.*]] = load <256 x i32>, <256 x i32>* [[IN:%.*]], align 64
 ; CHECK-NEXT:    [[TMP0:%.*]] = bitcast <256 x i32>* [[IN]] to i8*
 ; CHECK-NEXT:    [[TMP1:%.*]] = call x86_amx @llvm.x86.tileloadd64.internal(i16 [[M:%.*]], i16 [[N:%.*]], i8* [[TMP0]], i64 64)
-; CHECK-NEXT:    call void @llvm.x86.tilestored64.internal(i16 [[M]], i16 [[N]], i8* [[BUF:%.*]], i64 [[S:%.*]], x86_amx [[TMP1]]) [[ATTR3:#.*]]
+; CHECK-NEXT:    call void @llvm.x86.tilestored64.internal(i16 [[M]], i16 [[N]], i8* [[BUF:%.*]], i64 [[S:%.*]], x86_amx [[TMP1]]) [[ATTR3]]
 ; CHECK-NEXT:    ret <256 x i32> [[T1]]
 ;
 entry:
   %t1 = load <256 x i32>, <256 x i32>* %in, align 64
   %t2 = bitcast <256 x i32> %t1 to x86_amx
   call void @llvm.x86.tilestored64.internal(i16 %m, i16 %n, i8* %buf, i64 %s, x86_amx %t2) #3
   ret <256 x i32> %t1
 }