public JittedKernel(JitContext ctx)
{
_cfg = ctx.Cfg.Clone();
_ptx = ctx.Ptx;
_hir = ctx.Hir.Fluent(blk => blk.FreezeForever());
var flds = ctx.Allocator.Fields;
_memcpyHostToDevice = kernel =>
{
var args = new KernelArguments();
flds.Keys.ForEach(fld =>
{
var value = fld.GetValue(kernel);
if (flds[fld] is SlotLayout)
{
args.Add(value.In());
}
else if (flds[fld] is ArrayLayout)
{
var arr = value.AssertCast<Array>();
args.Add(arr.InOut());
var rank = arr.GetType().GetArrayRank();
0.UpTo(rank - 1).ForEach(i => args.Add(arr.GetLength(i).In()));
}
else
{
throw AssertionHelper.Fail();
}
});
return args;
};
_memcpyDeviceToHost = (result, kernel) =>
{
var idx = 0;
flds.Keys.ForEach(fld =>
{
Object value;
if (flds[fld] is SlotLayout)
{
value = result[idx];
idx += 1;
}
else if (flds[fld] is ArrayLayout)
{
value = result[idx];
idx += 3;
}
else
{
throw AssertionHelper.Fail();
}
fld.SetValue(kernel, value);
});
};
}
[DebuggerNonUserCode] public BlockExpander(Block source, ExpansionContext ctx)
{
Source = source;
Ctx = ctx;
Stmts = new Block();
ctx.Scope = Stmts;
var cloned_locals = source.Locals.Select(l => l.DeepClone());
cloned_locals.ForEach(local => ctx.Scope.Locals.Add(local));
var is_root = Ctx.Parent == null || (Ctx.Parent.Stack.Count() != Ctx.Stack.Count());
if (is_root) RetLabel = new Label();
else RetLabel.AssertNotNull();
source.ForEach(Expand);
if (is_root && Stmts.LastOrDefault() is Goto) Stmts.RemoveLast();
var gotos = Stmts.Family().OfType<Goto>().Where(@goto => @goto.LabelId == RetLabel.Id).ToReadOnly();
if (is_root && gotos.IsNotEmpty()) Stmts.Add(RetLabel);
}
protected internal override void TraverseBlock(Block block)
{
var prefixes = block.ToDictionary(s => s, s => new List<Action<IndentedWriter>>());
var evalOrder = block.CSharpEvaluationOrder();
var allRefs = evalOrder.OfType<Ref>().Where(
@ref => @ref.Sym != null && @ref.Sym.IsLocal()).ToReadOnly();
var r_scopes = allRefs.ToDictionary(@ref => @ref, @ref => @ref.Scope());
block.Locals.ForEach(loc =>
{
var refs = allRefs.Where(@ref => @ref.Sym == loc).ToReadOnly();
if (refs.IsEmpty()) return;
var firstUsage = refs.First();
var firstAss = firstUsage.Parent.AssertCast<Assign>().AssertNotNull();
(firstAss.Lhs is Ref && ((Ref)firstAss.Lhs).Sym == loc).AssertTrue();
if (ReferenceEquals(r_scopes[firstUsage], block))
{
prefixes[firstAss].Add(w => w.Write("var "));
}
else
{
var stmt = firstUsage.Hierarchy().SkipWhile(n => n.Parent != block).First();
prefixes[stmt].Add(w => w.WriteLine(String.Format("{0} {1};",
loc.Type == null ? "?" : loc.Type.GetCSharpRef(ToCSharpOptions.Informative), loc.Name)));
}
});
_writer.WriteLine("{");
_writer.Indent++;
block.ForEach(c =>
{
prefixes[c].ForEach(action => action(_writer));
Traverse(c);
if (IsOneliner(c))
{
if (c is Label) _writer.WriteLine(":");
else _writer.WriteLine(";");
}
});
_writer.Indent--;
_writer.WriteLine("}");
}
private JitCompiler(CpuConfig config, Type t_kernel, TypeBuilder t_xformed)
{
_config = config;
_t_kernel = t_kernel;
_m_kernel = _t_kernel.GetMethod("RunKernel", BF.All);
ValidateInputParameters();
// todo. think how to support multiple methods: *Dims, *Idxes, synchronization points
// todo. I'd even say that we should embed constants instead of Dims and cache such bodies between calls
// todo. if some dimension is equal to 1, then strip off corresponding loop
var lam = _m_kernel.Decompile();
_hir = lam.Body;
_params = lam.Sig.Syms;
_locals = lam.Body.LocalsRecursive();
_xhir = new Block();
_tids.Add("x", new Local("tid_x", typeof(int)));
_tids.Add("y", new Local("tid_y", typeof(int)));
_tids.Add("z", new Local("tid_z", typeof(int)));
_xhir.Locals.AddElements(_tids.Values);
InferLocalAllocationHints();
ReplicatePrivatelyAllocatedLocals();
LiftLocallyAllocatedLocals();
HoistGloballyAllocatedLocals();
_callsToSyncThreads = _hir.Family().OfType<Eval>().Where(eval =>
{
var m1 = eval.InvokedMethod();
var syncapi_syncThreads = typeof(ISyncApi).GetMethod("SyncThreads");
return m1.Api() == syncapi_syncThreads;
}).ToReadOnly();
TransformBlock(_hir, _xhir, false);
// todo. currently we don't support debuggable IL
// the latter implies correct PDB-mappings and working watch for replicated vars
_config.EmitDebuggableIL.AssertFalse();
_t_xformed = t_xformed;
_t_xformed.AddInterfaceImplementation(typeof(IBlockRunner));
_m_xformed = _t_xformed.DefineMethod("RunBlock", MA.Public, typeof(void), typeof(int3).MkArray());
_m_xformed.DefineParameter(1, ParmA.None, "blockIdx");
CompileTransformedHir();
}
public static Block DoRestoreTypeIs(Block hir)
{
if (hir.Family().OfType<TypeAs>().None()) return hir;
else
{
return hir.Transform((Operator op) =>
{
var is_rel = op.OperatorType.IsRelational();
if (is_rel)
{
var bin = op.AssertCast<BinaryOperator>();
var opt = bin.OperatorType;
var needs_xform = bin.Lhs is TypeAs ^ bin.Rhs is TypeAs;
if (needs_xform)
{
var type_as = bin.Lhs as TypeAs ?? bin.Rhs as TypeAs;
var other = bin.Lhs == type_as ? bin.Rhs : bin.Lhs;
other.AssertCast<Const>().AssertThat(c => c.Value == null);
var pos = opt == OperatorType.NotEqual || opt == OperatorType.GreaterThan;
var neg = opt == OperatorType.Equal || opt == OperatorType.LessThanOrEqual;
(pos || neg).AssertTrue();
var type_is = new TypeIs(type_as.Type, type_as.Target) as Expression;
return pos ? type_is : Operator.Negate(type_is);
}
else
{
return op.DefaultTransform();
}
}
else
{
return op.DefaultTransform();
}
}).AssertCast<Block>();
}
}
protected override void TraverseBlock(Block block)
{
block.ForEach(stmt => _ptx.emit(stmt));
}
protected internal override void TraverseBlock(Block block)
{
Dispatch(block);
}
public static Block DoSimplifyConditions(Block hir)
{
return hir.SimplifyConditions().AssertCast<Block>();
}
protected internal override Node TransformBlock(Block block)
{
return Dispatch(block);
}
public static Block DoRestoreUsings(Block hir)
{
// todo. to be implemented
return hir;
}
private void TransformBlock(Block hir, Block xhir, bool insideThreadLoop)
{
if (hir == null) return;
var deepFission = new List<Node>();
var regions = new List<ReadOnlyCollection<Node>>();
var curr_region = new List<Node>();
foreach (var stmt in hir)
{
if (_callsToSyncThreads.Contains(stmt as Eval))
{
if (curr_region.IsNotEmpty()) regions.Add(curr_region.ToReadOnly());
curr_region = new List<Node>();
}
else if (_callsToSyncThreads.Any(c => Set.Intersect(c.Hierarchy(), stmt.MkArray()).IsNotEmpty()))
{
if (curr_region.IsNotEmpty()) regions.Add(curr_region.ToReadOnly());
curr_region = new List<Node>();
deepFission.Add(stmt);
regions.Add(stmt.MkArray().ToReadOnly());
}
else
{
curr_region.Add(stmt);
}
}
if (curr_region.IsNotEmpty()) regions.Add(curr_region.ToReadOnly());
curr_region = null;
var privateLocals = hir.Locals.Where(l => Alloc(l) == allocPrivate).ToReadOnly();
var privateUsages = privateLocals.ToDictionary(l => l,
l => regions.Where(r => r.UsagesOfLocal(l).IsNotEmpty()).ToReadOnly());
privateUsages.ForEach(kvp =>
{
var needsReplication = kvp.Value.Count() > 1;
_needsReplication[kvp.Key] = needsReplication;
if (needsReplication)
{
_xhir.Insert(_lastIndex++, _replicatedInits[kvp.Key]);
_xhir.Locals.Add(_replicatedLocals[kvp.Key]);
}
});
foreach (var region in regions)
{
var xregion = new Block();
var needsToBeWrapped = !deepFission.Contains(region.SingleOrDefault2());
var regionIsInsideThreadLoop = insideThreadLoop || needsToBeWrapped;
foreach (var stmt in region)
{
if (stmt is Expression)
{
TransformExpression(((Expression)stmt), xregion, regionIsInsideThreadLoop);
}
else if (stmt is If)
{
TransformIf(((If)stmt), xregion, regionIsInsideThreadLoop);
}
else if (stmt is Loop)
{
TransformLoop(((Loop)stmt), xregion, regionIsInsideThreadLoop);
}
else
{
throw AssertionHelper.Fail();
}
}
if (needsToBeWrapped && !insideThreadLoop)
{
xhir.Add(new Loop
{
Init = new Block(new Assign(new Ref(_tids["z"]), new Const(0))),
Test = Operator.LessThan(new Ref(_tids["z"]), new Fld(typeof(int3).GetField("Z"), new Prop(typeof(IGridApi).GetProperty("BlockDim"), new Ref(_this), true))),
Body = new Block(new Loop
{
Init = new Block(new Assign(new Ref(_tids["y"]), new Const(0))),
Test = Operator.LessThan(new Ref(_tids["y"]), new Fld(typeof(int3).GetField("Y"), new Prop(typeof(IGridApi).GetProperty("BlockDim"), new Ref(_this), true))),
Body = new Block(new Loop
{
Init = new Block(new Assign(new Ref(_tids["x"]), new Const(0))),
Test = Operator.LessThan(new Ref(_tids["x"]), new Fld(typeof(int3).GetField("X"), new Prop(typeof(IGridApi).GetProperty("BlockDim"), new Ref(_this), true))),
Body = new Block(xregion.Children),
Iter = new Block(Operator.PreIncrement(new Ref(_tids["x"]))),
}),
Iter = new Block(Operator.PreIncrement(new Ref(_tids["y"]))),
}),
Iter = new Block(Operator.PreIncrement(new Ref(_tids["z"]))),
});
}
else
{
xhir.AddElements(xregion.Children);
}
}
var locals = hir.Locals.Except(_liftedLocals).ToReadOnly();
locals.ForEach(l => xhir.Locals.Add(l.DeepClone()));
}
private void Expand(Node node)
{
if (node is Expression)
{
var expr = (Expression)node;
var inlined = expr.Expand(Ctx);
inlined.Stmts.ForEach(Stmts.Add);
if (inlined.Result != null) Stmts.Add(inlined.Result);
}
else if (node is Block)
{
var block = (Block)node;
Stmts.Add(block.Expand(Ctx.SpinOff()));
}
else if (node is Break)
{
Stmts.Add(node);
}
else if (node is Continue)
{
Stmts.Add(node);
}
else if (node is Goto)
{
Stmts.Add(node);
}
else if (node is Label)
{
Stmts.Add(node);
}
else if (node is If)
{
var @if = (If)node;
var test = @if.Test.Expand(Ctx);
test.Stmts.ForEach(Stmts.Add);
test.Result.AssertNotNull();
var if_true = @if.IfTrue.Expand(Ctx.SpinOff());
var if_false = @if.IfFalse.Expand(Ctx.SpinOff());
var expanded = new If(test.Result, if_true, if_false);
Stmts.Add(expanded);
}
else if (node is Loop)
{
var loop = (Loop)node;
var test = loop.Test.Expand(Ctx);
test.Result.AssertNotNull();
var init = loop.Init.Expand(Ctx.SpinOff());
var iter = loop.Iter.Expand(Ctx.SpinOff());
var body = loop.Body.Expand(Ctx.SpinOff());
var prepend_test = loop.IsWhileDo && test.Stmts.IsNotEmpty();
if (init.IsNotEmpty() && prepend_test) { Stmts.Add(init); init = new Block(); }
if (prepend_test) test.Stmts.ForEach(Stmts.Add);
test.Stmts.ForEach(iter.Add);
var xloop = new Loop(test.Result, body, loop.IsWhileDo){Init = init, Iter = iter};
var cloned_locals = loop.Locals.Select(l => l.DeepClone());
cloned_locals.ForEach(local => xloop.Locals.Add(local));
Stmts.Add(xloop);
}
else if (node is Return)
{
var ret = (Return)node;
(ret.Value == null).AssertEquiv(Ret == null);
if (ret.Value != null) Expand(new Assign(Ret, ret.Value));
Stmts.Add(new Goto(RetLabel));
}
else if (node is Try || node is Clause || node is Throw ||
node is Using || node is Iter)
{
// todo. implement support for non-linear control flow
// this is only possible when we fully implement decompilation of tries
// until now I leave this marked as "to be implemented"
throw AssertionHelper.Fail();
}
else
{
throw AssertionHelper.Fail();
}
}
[DebuggerNonUserCode] public static Block ExpandBlock(Block source, ExpansionContext ctx) { return new BlockExpander(source, ctx).Stmts; }
public static Block DoRestoreBooleans(Block hir)
{
// todo. also think about how to support nullables
var types = hir.InferTypes();
return (Block)hir.Transform(
(If @if) =>
{
var is_bool = types[@if.Test] == typeof(bool);
var ensuredTest = is_bool ? @if.Test.CurrentTransform() : @if.Test.EnsureBoolean(types);
var ifTrue = @if.IfTrue.CurrentTransform();
var ifFalse = @if.IfFalse.CurrentTransform();
return new If(ensuredTest, ifTrue, ifFalse);
},
(Assign ass) =>
{
if (types[ass.Lhs] == typeof(bool))
{
var lhs = ass.Lhs.CurrentTransform();
var ensuredRhs = ass.Rhs.EnsureBoolean(types);
return new Assign(lhs, ensuredRhs);
}
else
{
return ass.DefaultTransform();
}
},
(Apply app) =>
{
ReadOnlyCollection<Type> targs = null;
if (app.Callee is Prop)
{
var p = app.Callee.InvokedProperty();
if (p != null) targs = p.GetIndexParameters().Select(pi => pi.ParameterType).ToReadOnly();
}
else
{
var ftype = types[app.Callee];
if (ftype != null) targs = ftype.DasmFType().SkipLast(1).ToReadOnly();
}
if (targs == null) return app.DefaultTransform();
var ensuredArgs = targs.Select((targ, i) =>
{
var arg = app.Args[i];
var is_bool = targ == typeof(bool);
return is_bool ? arg.EnsureBoolean(types) : arg.CurrentTransform();
});
var callee = app.Callee.CurrentTransform();
return new Apply(callee, ensuredArgs);
},
(Conditional cond) =>
{
var anyBool = types[cond.IfTrue] == typeof(bool) || types[cond.IfFalse] == typeof(bool);
var is_bool = types[cond.Test] == typeof(bool);
var ensuredTest = is_bool ? cond.Test.CurrentTransform() : cond.Test.EnsureBoolean(types);
var ifTrue = anyBool ? cond.IfTrue.EnsureBoolean(types) : cond.IfTrue.CurrentTransform();
var ifFalse = anyBool ? cond.IfFalse.EnsureBoolean(types) : cond.IfFalse.CurrentTransform();
return new Conditional(ensuredTest, ifTrue, ifFalse);
},
(Operator op) =>
{
var opt = op.OperatorType;
var anyBool = op.Args.Any(arg => types[arg] == typeof(bool));
var alwaysBool = opt == OperatorType.AndAlso ||
opt == OperatorType.OrElse || opt == OperatorType.Not;
var ensureBool = anyBool || alwaysBool;
if (ensureBool)
{
if (opt == OperatorType.And) opt = OperatorType.AndAlso;
if (opt == OperatorType.Or) opt = OperatorType.OrElse;
if (opt == OperatorType.Negate) opt = OperatorType.Not;
var ensuredArgs = op.Args.Select(arg => arg.EnsureBoolean(types));
return Operator.Create(opt, ensuredArgs);
}
else
{
return op.DefaultTransform();
}
});
}
public static void DoInferLoopIters(Block hir)
{
foreach (var loop in hir.Family().OfType<Loop>())
{
if (loop.Iter.IsNotEmpty())
{
loop.Iter.RemoveElements(n => n == null);
}
else
{
if (loop.Init.IsNotEmpty())
{
Func<Expression, bool> isLvalueOfLool = e =>
{
if (e == null)
{
return false;
}
else if (e is Ref)
{
var @ref = e.AssertCast<Ref>();
var local = @ref == null ? null : @ref.Sym as Local;
return local != null && loop.Locals.Contains(local);
}
else if (e is Slot)
{
var slot = e.AssertCast<Slot>();
var @ref = slot == null ? null : slot.This as Ref;
var local = @ref == null ? null : @ref.Sym as Local;
return local != null && loop.Locals.Contains(local);
}
else if (e is Apply)
{
var app = e.AssertCast<Apply>();
var callee = app == null ? null : app.Callee;
var prop = callee as Prop;
var @ref = prop == null ? null : prop.This as Ref;
var local = @ref == null ? null : @ref.Sym as Local;
return local != null && loop.Locals.Contains(local);
}
else
{
return false;
}
};
var lil = loop.Body.LastOrDefault() as Expression;
var lilIsIter = false;
if (lil is Assign)
{
var ass = lil.AssertCast<Assign>();
lilIsIter = isLvalueOfLool(ass.Lhs);
}
else if (lil is Operator)
{
var op = lil.AssertCast<Operator>();
if (op.OperatorType.IsAssign())
{
lilIsIter = isLvalueOfLool(op.Args.First());
}
}
else if (lil is Eval)
{
var m = lil.InvokedMethod();
if (m != null)
{
var args = lil.InvocationArgs();
if (m.IsInstance() || m.IsExtension())
{
lilIsIter = isLvalueOfLool(args.FirstOrDefault());
}
else
{
lilIsIter = args.Any(isLvalueOfLool);
}
}
}
if (lilIsIter)
{
loop.Iter = new Block();
loop.Iter.Add(lil);
loop.Body.Remove(lil);
}
}
}
}
}
// todo. verify that init and iter can live outside the thread loop (when insideThreadLoop is false)
private void TransformLoop(Loop loop, Block xregion, bool insideThreadLoop)
{
var x_init = new Block(); TransformBlock(loop.Init, x_init, true);
var x_test = TransformExpression(loop.Test, insideThreadLoop);
var x_body = new Block(); TransformBlock(loop.Body, x_body, insideThreadLoop);
var x_iter = new Block(); TransformBlock(loop.Iter, x_iter, true);
var x_loop = new Loop(x_test, x_body);
x_loop.Locals.SetElements(loop.Locals);
x_loop.Init = new Block(x_init.Children);
x_loop.Iter = new Block(x_iter.Children);
xregion.Add(x_loop);
}
private void TransformIf(If @if, Block xregion, bool insideThreadLoop)
{
var x_test = TransformExpression(@if.Test, insideThreadLoop);
var x_iftrue = new Block(); TransformBlock(@if.IfTrue, x_iftrue, insideThreadLoop);
var x_iffalse = new Block(); TransformBlock(@if.IfFalse, x_iffalse, insideThreadLoop);
xregion.Add(new If(x_test, x_iftrue, x_iffalse));
}
private void TransformExpression(Expression e, Block xregion, bool insideThreadLoop)
{
var xe = TransformExpression(e, insideThreadLoop);
if (xe != null) xregion.Add(xe);
}
protected override void TraverseBlock(Block block)
{
Func<Type, Type> safeType = t => t.IsRectMdArray() ? typeof(Object) : t;
block.Locals.ForEach(l => locals.Add(l, il.DeclareLocal(safeType(l.Type.AssertNotNull()))));
block.ForEach(Traverse);
}
protected internal override void TraverseBlock(Block block)
{
block.ForEach(Traverse);
Types.Add(block, null);
}
protected internal virtual void TraverseBlock(Block block) { block.Unsupported(); }
protected internal virtual Node TransformBlock(Block block) { return block.AcceptTransformer(this, true); }
