[AIE2P] legalizer support for G_FMUL

Xilinx · Feb 20, 2025 · 02cc673 · 02cc673
1 parent 69ca418
commit 02cc673
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 2 deletions.
diff --git a/llvm/lib/Target/AIE/AIELegalizerHelper.cpp b/llvm/lib/Target/AIE/AIELegalizerHelper.cpp
@@ -27,6 +27,7 @@
 #include "llvm/IR/IntrinsicsAIE2.h"
 #include "llvm/IR/IntrinsicsAIE2P.h"
 #include "llvm/Support/ErrorHandling.h"
+#include <cassert>
 
 namespace llvm {
 
@@ -1157,6 +1158,51 @@ bool AIELegalizerHelper::legalizeG_FABS(LegalizerHelper &Helper,
   return true;
 }
 
+bool AIELegalizerHelper::legalizeG_FMUL(LegalizerHelper &Helper,
+                                        MachineInstr &MI) const {
+  assert(ST.isAIE2P() && "Custom legalization supported for AIE2P only");
+
+  MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
+  MachineRegisterInfo &MRI = *MIRBuilder.getMRI();
+
+  const Register DstReg = MI.getOperand(0).getReg();
+  assert(MRI.getType(DstReg) == LLT::scalar(16) &&
+         "Expected bfloat16 type in custom legalization.");
+
+  Register SrcLHS = MI.getOperand(1).getReg();
+  Register SrcRHS = MI.getOperand(2).getReg();
+
+  const LLT InsertVecLLT = V32BF16;
+  const unsigned InsertEltOpc =
+      ST.getInstrInfo()->getGenericInsertVectorEltOpcode();
+
+  const Register IdxReg = MIRBuilder.buildConstant(S32, 0).getReg(0);
+  const Register UndefVec512 = MIRBuilder.buildUndef(InsertVecLLT).getReg(0);
+
+  SrcLHS = MIRBuilder
+               .buildInstr(InsertEltOpc, {InsertVecLLT},
+                           {UndefVec512, SrcLHS, IdxReg})
+               .getReg(0);
+  SrcRHS = MIRBuilder
+               .buildInstr(InsertEltOpc, {InsertVecLLT},
+                           {UndefVec512, SrcRHS, IdxReg})
+               .getReg(0);
+
+  Register Res =
+      MIRBuilder.buildInstr(MI.getOpcode(), {V32BF16}, {SrcLHS, SrcRHS})
+          .getReg(0);
+
+  const unsigned ExtractEltOpc =
+      ST.getInstrInfo()->getGenericExtractVectorEltOpcode(/*SignExt*/ true);
+  Res = MIRBuilder.buildInstr(ExtractEltOpc, {S32}, {Res, IdxReg}).getReg(0);
+  Res = MIRBuilder.buildAssertInstr(TargetOpcode::G_ASSERT_SEXT, {S32}, Res, 16)
+            .getReg(0);
+  MIRBuilder.buildTrunc(DstReg, Res);
+
+  MI.eraseFromParent();
+  return true;
+}
+
 bool AIELegalizerHelper::legalizeG_FADD_G_FSUB(LegalizerHelper &Helper,
                                                MachineInstr &MI) const {
   MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;

diff --git a/llvm/lib/Target/AIE/AIELegalizerHelper.h b/llvm/lib/Target/AIE/AIELegalizerHelper.h
@@ -38,7 +38,9 @@ class AIELegalizerHelper {
   const LLT V32BF16 = LLT::fixed_vector(32, 16);
   const LLT V32FP32 = LLT::fixed_vector(32, 32);
   const LLT V32ACC32 = LLT::fixed_vector(32, 32);
+  const LLT V64BF16 = LLT::fixed_vector(64, 16);
   const LLT V64FP32 = LLT::fixed_vector(64, 32);
+  const LLT V128BF16 = LLT::fixed_vector(128, 16);
 
 public:
   AIELegalizerHelper(const AIEBaseSubtarget &ST);
@@ -70,6 +72,7 @@ class AIELegalizerHelper {
   bool legalizeG_FPEXT(LegalizerHelper &Helper, MachineInstr &MI) const;
   bool legalizeG_FABS(LegalizerHelper &Helper, MachineInstr &MI) const;
   bool legalizeG_FADD_G_FSUB(LegalizerHelper &Helper, MachineInstr &MI) const;
+  bool legalizeG_FMUL(LegalizerHelper &Helper, MachineInstr &MI) const;
   bool legalizeG_SELECT(LegalizerHelper &Helper, MachineInstr &MI,
                         const unsigned MaxBitSize = 512) const;
   bool legalizeG_BITCAST(LegalizerHelper &Helper, MachineInstr &MI) const;

diff --git a/llvm/lib/Target/AIE/aie2p/AIE2PLegalizerInfo.cpp b/llvm/lib/Target/AIE/aie2p/AIE2PLegalizerInfo.cpp
@@ -225,12 +225,17 @@ AIE2PLegalizerInfo::AIE2PLegalizerInfo(const AIE2PSubtarget &ST)
 
   getActionDefinitionsBuilder(G_FABS).customFor({S16, S32, S64}).scalarize(0);
 
+  getActionDefinitionsBuilder(G_FMUL)
+      .legalFor({V64S16, V32S16})
+      .customFor({S16})
+      .libcallFor({S32, S64});
+
   getActionDefinitionsBuilder({G_FADD, G_FSUB})
       .legalFor({AccV64S32})
       .customFor({S16})
       .libcallFor({S32, S64});
 
-  getActionDefinitionsBuilder({G_FMUL, G_FDIV, G_FREM})
+  getActionDefinitionsBuilder({G_FDIV, G_FREM})
       .clampScalar(0, S32, S64)
       .libcallFor({S32, S64});
 
@@ -723,6 +728,8 @@ bool AIE2PLegalizerInfo::legalizeCustom(
   case TargetOpcode::G_FADD:
   case TargetOpcode::G_FSUB:
     return AIEHelper.legalizeG_FADD_G_FSUB(Helper, MI);
+  case TargetOpcode::G_FMUL:
+    return AIEHelper.legalizeG_FMUL(Helper, MI);
   case TargetOpcode::G_BUILD_VECTOR:
     return AIEHelper.legalizeG_BUILD_VECTOR(Helper, MI);
   case TargetOpcode::G_UNMERGE_VALUES:

diff --git a/llvm/test/CodeGen/AIE/GlobalISel/legalize-fmul.mir b/llvm/test/CodeGen/AIE/GlobalISel/legalize-fmul.mir
@@ -5,7 +5,6 @@
 # (c) Copyright 2024 Advanced Micro Devices, Inc. or its affiliates
 
 # RUN: llc -mtriple aie2 -run-pass=legalizer %s -verify-machineinstrs -o - | FileCheck -DVER=2 --check-prefix=COMMON --check-prefix=AIE2 %s
-# RUN: llc -mtriple aie2p -run-pass=legalizer %s -verify-machineinstrs -o - | FileCheck -DVER=2p --check-prefix=COMMON --check-prefix=AIE2P %s
 
 ---
 name: test_fmul_bfloat16