def test_dynamic_shape(self):
with kernel_arena_scope():
dN = te.VarHandle("n", te.Dtype.Int)
A = te.Placeholder('A', te.Dtype.Double, [dN])
B = te.Placeholder('B', te.Dtype.Double, [dN])
def compute(i):
return A.load([i]) - B.load([i])
C = te.Compute('C', [te.DimArg(dN, 'i')], compute)
loopnest = te.LoopNest([C])
loopnest.prepare_for_codegen()
stmt = te.simplify(loopnest.root_stmt())
cg = te.construct_codegen('ir_eval', stmt, [A, B, C, dN])
def test_with_shape(n):
tA = torch.randn(n, dtype=torch.double)
tB = torch.randn(n, dtype=torch.double)
tC = torch.empty(n, dtype=torch.double)
cg.call([tA, tB, tC, n])
torch.testing.assert_allclose(tA - tB, tC)
test_with_shape(8)
test_with_shape(31)
def construct_adder(n: int, dtype=te.Dtype.Float):
dN = te.ExprHandle.int(n)
A = te.Placeholder('A', dtype, [dN])
B = te.Placeholder('B', dtype, [dN])
def compute(i):
return A.load([i]) + B.load([i])
C = te.Compute('C', [te.DimArg(dN, 'i')], compute)
loopnest = te.LoopNest([C])
loopnest.prepare_for_codegen()
stmt = te.simplify(loopnest.root_stmt())
return te.construct_codegen('ir_eval', stmt, [A, B, C])
def get_dim_args(dims):
dim_args = []
for dim in dims:
dim_args.append(
te.DimArg(te.ExprHandle.int(dim), 'i' + str(len(dim_args))))
return dim_args
def get_dim_args(dims):
dim_args = []
for dim in dims:
dim_args.append(te.DimArg(dim, 'i' + str(len(dim_args))))
return dim_args