[AIE2P] legalizer support for G_FMUL

Xilinx · Feb 20, 2025 · 0416372 · 0416372
1 parent d9dec55
commit 0416372
Show file tree

Hide file tree

Showing 5 changed files with 135 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
@@ -70,6 +70,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

diff --git a/llvm/test/CodeGen/AIE/aie2p/GlobalIsel/legalize-fmul.mir b/llvm/test/CodeGen/AIE/aie2p/GlobalIsel/legalize-fmul.mir
@@ -0,0 +1,80 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 4
+# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#
+# (c) Copyright 2024 Advanced Micro Devices, Inc. or its affiliates
+
+# RUN: llc -mtriple aie2p -run-pass=legalizer %s -verify-machineinstrs -o - | FileCheck %s
+
+---
+name: test_fmul_s16
+body: |
+  bb.0:
+    liveins: $r1, $r2
+    ; CHECK-LABEL: name: test_fmul_s16
+    ; CHECK: liveins: $r1, $r2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY $r1
+    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:_(s16) = G_TRUNC [[COPY]](s32)
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(s32) = COPY $r2
+    ; CHECK-NEXT: [[TRUNC1:%[0-9]+]]:_(s16) = G_TRUNC [[COPY1]](s32)
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(<32 x s16>) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[AIE_INSERT_VECTOR_ELT:%[0-9]+]]:_(<32 x s16>) = G_AIE_INSERT_VECTOR_ELT [[DEF]], [[TRUNC]](s16), [[C]](s32)
+    ; CHECK-NEXT: [[AIE_INSERT_VECTOR_ELT1:%[0-9]+]]:_(<32 x s16>) = G_AIE_INSERT_VECTOR_ELT [[DEF]], [[TRUNC1]](s16), [[C]](s32)
+    ; CHECK-NEXT: [[FMUL:%[0-9]+]]:_(<32 x s16>) = G_FMUL [[AIE_INSERT_VECTOR_ELT]], [[AIE_INSERT_VECTOR_ELT1]]
+    ; CHECK-NEXT: [[AIE_SEXT_EXTRACT_VECTOR_ELT:%[0-9]+]]:_(s32) = G_AIE_SEXT_EXTRACT_VECTOR_ELT [[FMUL]](<32 x s16>), [[C]](s32)
+    ; CHECK-NEXT: [[ASSERT_SEXT:%[0-9]+]]:_(s32) = G_ASSERT_SEXT [[AIE_SEXT_EXTRACT_VECTOR_ELT]], 16
+    ; CHECK-NEXT: $r0 = COPY [[ASSERT_SEXT]](s32)
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit $r0
+    %0:_(s32) = COPY $r1
+    %1:_(s16) = G_TRUNC %0(s32)
+    %2:_(s32) = COPY $r2
+    %3:_(s16) = G_TRUNC %2(s32)
+    %4:_(s16) = G_FMUL %1, %3
+    %5:_(s32) = G_ANYEXT %4(s16)
+    $r0 = COPY %5(s32)
+    PseudoRET implicit $lr, implicit $r0
+...
+
+---
+name: test_fmul_vec_1024
+body: |
+  bb.0:
+    liveins: $dm0, $dm1
+    ; CHECK-LABEL: name: test_fmul_vec_1024
+    ; CHECK: liveins: $dm0, $dm1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<64 x s16>) = COPY $cml0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(<64 x s16>) = COPY $cml1
+    ; CHECK-NEXT: [[FMUL:%[0-9]+]]:_(<64 x s16>) = G_FMUL [[COPY]], [[COPY1]]
+    ; CHECK-NEXT: $cml0 = COPY [[FMUL]](<64 x s16>)
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit $cml0
+    %0:_(<64 x s16>) = COPY $cml0
+    %1:_(<64 x s16>) = COPY $cml1
+    %2:_(<64 x s16>) = G_FMUL %0, %1
+    $cml0 = COPY %2(<64 x s16>)
+    PseudoRET implicit $lr, implicit $cml0
+...
+
+---
+name: test_fmul_vec_512
+body: |
+  bb.0:
+    liveins: $dm0, $dm1
+    ; CHECK-LABEL: name: test_fmul_vec_512
+    ; CHECK: liveins: $dm0, $dm1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<32 x s16>) = COPY $bmll0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(<32 x s16>) = COPY $bmll1
+    ; CHECK-NEXT: [[FMUL:%[0-9]+]]:_(<32 x s16>) = G_FMUL [[COPY]], [[COPY1]]
+    ; CHECK-NEXT: $bmll0 = COPY [[FMUL]](<32 x s16>)
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit $bmll0
+    %0:_(<32 x s16>) = COPY $bmll0
+    %1:_(<32 x s16>) = COPY $bmll1
+    %2:_(<32 x s16>) = G_FMUL %0, %1
+    $bmll0 = COPY %2(<32 x s16>)
+    PseudoRET implicit $lr, implicit $bmll0
+...
+