跳到主要内容

CASK IR

Cask IR

The representation that we work on to generate a fused gemm + epilogue call. This IR interface is exposed by the CASK library, and we translate a subgraph in Myelin to instructions in this IR, and then codegen this IR to get a GraphShader.

Key Concepts:

  • Tensor: multi-dimentional arrays of elements
  • Element
  • Shader: function of tensor variables
    • Backbone shader: one and only one needed to lowering into GraphShader
    • ReducePerTile: local reduction
    • Reduce: global reduction
  • Element-wise functions
  • Tensorize
    • rank: dimension of output
    • dims: use which dimensions of the input

实例

def foo(x : float, y : float) -> { x + y; }  
rank = 3 // dimension of output is 3, [B, M, N]
dims = [[1], // dimension of tensor X is 1, use M only
[0, 2, 1]] // dimension of tensor Y is 3, use B, M, N

def bar = tensorize(foo, 3, dims)


bar (X : float[M], Y : float[B, N, M]) -> float[B, M, N] {
float r[M, N, K];
for (i, j, k) in [B, M, N] {
r[i, j, k] = foo(X[j], Y[i, k, j]);
}
return r;

相关内容

参考文献