Skip to content

Commit

Permalink
Fixed PTXCodeGenerator to emit individual target bindings.
Browse files Browse the repository at this point in the history
  • Loading branch information
@m4rs-mt
m4rs-mt committed Sep 5, 2023
1 parent 09c262f commit ec4e7d8
Show file tree
Hide file tree
Showing 2 changed files with 151 additions and 52 deletions.
89 changes: 88 additions & 1 deletion Src/ILGPU/Backends/PTX/PTXCodeGenerator.Terminators.cs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
// ---------------------------------------------------------------------------------------
// ILGPU
// Copyright (c) 2018-2022 ILGPU Project
// Copyright (c) 2018-2023 ILGPU Project
// www.ilgpu.net
//
// File: PTXCodeGenerator.Terminators.cs
Expand All @@ -9,8 +9,10 @@
// Source License. See LICENSE.txt for details.
// ---------------------------------------------------------------------------------------

using ILGPU.IR;
using ILGPU.IR.Values;
using ILGPU.Util;
using System.Diagnostics.CodeAnalysis;

namespace ILGPU.Backends.PTX
{
Expand All @@ -24,15 +26,50 @@ public void GenerateCode(ReturnTerminator returnTerminator)
var resultRegister = Load(returnTerminator.ReturnValue);
EmitStoreParam(ReturnParamName, resultRegister);
}

Command(
Uniforms.IsUniform(returnTerminator)
? PTXInstructions.UniformReturnOperation
: PTXInstructions.ReturnOperation);
}

private bool NeedSeparatePhiBindings(
BasicBlock basicBlock,
BasicBlock target,
[NotNullWhen(true)]
out PhiBindings.PhiBindingCollection bindings)
{
if (!phiBindings.TryGetBindings(target, out bindings))
return false;

// Check whether there are bindings pointing to different blocks
foreach (var (phiValue, _) in bindings)
{
if (phiValue.BasicBlock != target)
return true;
}

// We were not able to find misleading data
return false;
}

/// <summary>
/// Generates phi bindings for jumping to a specific target block.
/// </summary>
/// <param name="current">The current block.</param>
private void GeneratePhiBindings(BasicBlock current)
{
if (!phiBindings.TryGetBindings(current, out var bindings))
return;
BindPhis(bindings, target: null);
}

/// <summary cref="IBackendCodeGenerator.GenerateCode(UnconditionalBranch)"/>
public void GenerateCode(UnconditionalBranch branch)
{
// Bind phis
GeneratePhiBindings(branch.BasicBlock);

if (Schedule.IsImplicitSuccessor(branch.BasicBlock, branch.Target))
return;

Expand Down Expand Up @@ -60,6 +97,53 @@ public void GenerateCode(IfBranch branch)
? PTXInstructions.UniformBranchOperation
: PTXInstructions.BranchOperation;

// Gather phi bindings and test both, true and false targets
if (phiBindings.TryGetBindings(branch.BasicBlock, out var bindings) &&
(bindings.NeedSeparateBindingsFor(trueTarget) ||
bindings.NeedSeparateBindingsFor(falseTarget)))
{
// We need to emit different bindings in each branch
if (branch.IsInverted)
Utilities.Swap(ref trueTarget, ref falseTarget);

// Declare a temporary jump target to skip true branches
var tempLabel = DeclareLabel();
using (var command = BeginCommand(
branchOperation,
new PredicateConfiguration(condition, isTrue: false)))
{
command.AppendLabel(tempLabel);
}

// Bind all true phis
BindPhis(bindings, trueTarget);

// Jump to true target in the current case
using (var command = BeginCommand(branchOperation))
{
var targetLabel = blockLookup[trueTarget];
command.AppendLabel(targetLabel);
}

// Mark the false case label and bind all values
MarkLabel(tempLabel);
BindPhis(bindings, falseTarget);

if (!Schedule.IsImplicitSuccessor(branch.BasicBlock, falseTarget))
{
// Jump to false target in the else case
using var command = BeginCommand(branchOperation);
var targetLabel = blockLookup[falseTarget];
command.AppendLabel(targetLabel);
}

// Skip further bindings an branches
return;
}

// Generate phi bindings for all blocks
BindPhis(bindings, target: null);

// The current schedule has inverted all if conditions with implicit branch
// targets to simplify the work of the PTX assembler
if (Schedule.IsImplicitSuccessor(branch.BasicBlock, trueTarget))
Expand Down Expand Up @@ -151,6 +235,9 @@ public void GenerateCode(SwitchBranch branch)
}
Builder.AppendLine(";");

// Generate all phi bindings for all cases
GeneratePhiBindings(branch.BasicBlock);

using (var command = BeginCommand(
isUniform
? PTXInstructions.UniformBranchIndexOperation
Expand Down
114 changes: 63 additions & 51 deletions Src/ILGPU/Backends/PTX/PTXCodeGenerator.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -317,10 +317,11 @@ protected static string GetParameterName(Parameter parameter) =>
#region Instance

private int labelCounter;
private readonly Dictionary<BasicBlock, string> blockLookup =
new Dictionary<BasicBlock, string>();
private readonly Dictionary<(Encoding, string), string> stringConstants =
new Dictionary<(Encoding, string), string>();
private readonly Dictionary<BasicBlock, string> blockLookup = new();

private readonly Dictionary<(Encoding, string), string> stringConstants = new();
private readonly PhiBindings phiBindings;
private readonly Dictionary<Value, Register> intermediatePhiRegisters;
private readonly string labelPrefix;

/// <summary>
Expand Down Expand Up @@ -352,6 +353,12 @@ internal PTXCodeGenerator(in GeneratorArgs args, Method method, Allocas allocas)
args.Properties.GetPTXBackendMode() == PTXBackendMode.Enhanced
? Method.Blocks.CreateOptimizedPTXSchedule()
: Method.Blocks.CreateDefaultPTXSchedule();

// Create phi bindings and initialize temporary phi registers
phiBindings = Schedule.ComputePhiBindings(
new PhiBindingAllocator(this));
intermediatePhiRegisters = new Dictionary<Value, Register>(
phiBindings.MaxNumIntermediatePhis);
}

#endregion
Expand Down Expand Up @@ -517,10 +524,6 @@ protected void GenerateCodeInternal(int registerOffset)
}

// Find all phi nodes, allocate target registers and setup internal mapping
var phiBindings = Schedule.ComputePhiBindings(
new PhiBindingAllocator(this));
var intermediatePhiRegisters = new Dictionary<Value, Register>(
phiBindings.MaxNumIntermediatePhis);
Builder.AppendLine();

// Generate code
Expand Down Expand Up @@ -555,59 +558,68 @@ protected void GenerateCodeInternal(int registerOffset)

DebugInfoGenerator.ResetLocation();

// Wire phi nodes
if (phiBindings.TryGetBindings(block, out var bindings))
{
// Assign all phi values
foreach (var (phiValue, value) in bindings)
{
// Load the current phi target register
var phiTargetRegister = Load(phiValue);

// Check for an intermediate phi value
if (bindings.IsIntermediate(phiValue))
{
var intermediateRegister = AllocateType(phiValue.Type);
intermediatePhiRegisters.Add(phiValue, intermediateRegister);

// Move this phi value into a temporary register for reuse
EmitComplexCommand(
PTXInstructions.MoveOperation,
new PhiMoveEmitter(),
intermediateRegister,
phiTargetRegister);
}

// Determine the source value from which we need to copy from
var sourceRegister = intermediatePhiRegisters
.TryGetValue(value, out var tempRegister)
? tempRegister
: Load(value);

// Move contents
EmitComplexCommand(
PTXInstructions.MoveOperation,
new PhiMoveEmitter(),
phiTargetRegister,
sourceRegister);
}

// Free temporary registers
foreach (var register in intermediatePhiRegisters.Values)
Free(register);
intermediatePhiRegisters.Clear();
}

// Build terminator
this.GenerateCodeFor(block.Terminator.AsNotNull());
Builder.AppendLine();

// Free temporary registers
foreach (var register in intermediatePhiRegisters.Values)
Free(register);
intermediatePhiRegisters.Clear();
}

// Finish function and append register information
Builder.AppendLine("}");
Builder.Insert(registerOffset, GenerateRegisterInformation("\t"));
}

/// <summary>
/// Binds all phi values of the current block flowing through an edge to the
/// target block.
/// </summary>
private void BindPhis(
PhiBindings.PhiBindingCollection bindings,
BasicBlock? target)
{
// Assign all phi values
foreach (var (phiValue, value) in bindings)
{
// Reject phis not flowing to the target edge
if (target is not null && phiValue.BasicBlock != target)
continue;

// Load the current phi target register
var phiTargetRegister = Load(phiValue);

// Check for an intermediate phi value
if (bindings.IsIntermediate(phiValue))
{
var intermediateRegister = AllocateType(phiValue.Type);
intermediatePhiRegisters.Add(phiValue, intermediateRegister);

// Move this phi value into a temporary register for reuse
EmitComplexCommand(
PTXInstructions.MoveOperation,
new PhiMoveEmitter(),
intermediateRegister,
phiTargetRegister);
}

// Determine the source value from which we need to copy from
var sourceRegister = intermediatePhiRegisters
.TryGetValue(value, out var tempRegister)
? tempRegister
: Load(value);

// Move contents
EmitComplexCommand(
PTXInstructions.MoveOperation,
new PhiMoveEmitter(),
phiTargetRegister,
sourceRegister);
}
}

/// <summary>
/// Setups local or shared allocations.
/// </summary>
Expand Down

0 comments on commit ec4e7d8

Please sign in to comment.