“Tensor LLVM Extensions Proposed For Targeting AI Accelerators, Emerging Hardware”
( via/context: Tensor LLVM Extensions Proposed For Targeting AI Accelerators, Emerging Hardware - Phoronix 14 November 2021 )
And a Julia example inside in the design document !!! ( big google doc )
Lowering Example Using Julia
Julia example using Knet.jl:
; Tensors in Knet are initialized using a reshape operation. The first argument to this function is a one-dimensional array with the data for the tensor (in this example, 1.0:6.0 means that the array is filled with values between 1 and 6, and the second argument is the shape of the tensor produced by the reshape function.w = reshape([1.0:6.0...], (1,1,3,2)) x = reshape([7.0:10.0...], (1,2,2)) y = conv4(w, x; padding=(1,2);stride=3)
The above convolution operation in Julia will be lowered down to LLVM IR to produce the following
%w = call token llvm.tensor.typeinfo(<6 x i32> <i32 1, i32 2, i32 3, i32 4, i32 5, i32 6>, <4 x i32> <i32 1, i32 1, i32 3, i32 2>, <4 x i32> <i32 0, i32 1, i32 2, i32 3>, <4 x i32> <i32 0, i32 0, i32 0, i32 0>) %x = call token llvm.tensor.typeinfo(<4 x i32> <i32 7, i32 8, i32 9, i32 10>, <3 x i32> <i32 1, i32 2, i32 2>, <3 x i32> <i32 0, i32 1, i32 2>, <3 x i32> <i32 0, i32 0, i32 0>) %pad_w = call <28 x i32> llvm.tensor.pad(token %w, i32 0, <3 x i32> <i32 0, i32 1, i32 2>, <3 x i32> <i32 0, i32 1, i32 2>, <3 x i32> <i32 0, i32 0, i32 0>) %padded_w = call token llvm.tensor.typeinfo(<28 x i32> %pad_w, <4 x i32> <i32 1, i32 1, i32 4, i32 4>, <4 x i32> <i32 0, i32 1, i32 2, i32 3>, <4 x i32> <i32 0, i32 0, i32 2, i32 4>) <vector_ty> llvm.tensor.convolution(token %padded_w, token %x, <4 x i32> <1, 1, 1, 3>, <4 x i32> <i32 0, i32 0, i32 0, i32 0>, <4 x i32> <i32 0, i32 0, i32 0, i32 0>)