samples/custom_dispatch/cpu/plugin/system_example.mlir

// This example demonstrates calling dynamically imported functions in the
// runtime. Alternatively the functions can be embedded into the compiled IREE
// programs for hermetic deployment (see custom_dispatch/cpu/embedded/).

// NOTE: this file is identical to system_example.mlir besides the lit config
// controlling the iree-run-module flag.
// TODO(benvanik): find a way to share the files (environment variables saying
// what types to run, etc).

// RUN: iree-compile --iree-hal-target-backends=llvm-cpu %s | \
// RUN: iree-run-module \
// RUN:     --device=local-sync \
// RUN:     --executable_plugin=$IREE_BINARY_DIR/samples/custom_dispatch/cpu/plugin/system_plugin$IREE_DYLIB_EXT \
// RUN:     --module=- \
// RUN:     --function=mixed_invocation \
// RUN:     --input=8xf32=2 \
// RUN:     --input=8xf32=4 | \
// RUN: FileCheck %s --check-prefix=CHECK-SYSTEM

// CHECK-SYSTEM: EXEC @mixed_invocation
// simple_mul_workgroup
// CHECK-SYSTEM: processor_id=
// CHECK-SYSTEM: processor_data[0]=
// CHECK-SYSTEM: mul[0:0](2 * 4 = 8)
// CHECK-SYSTEM: mul[0:1](2 * 4 = 8)
// CHECK-SYSTEM: mul[0:2](2 * 4 = 8)
// CHECK-SYSTEM: mul[0:3](2 * 4 = 8)
// CHECK-SYSTEM: mul[0:4](2 * 4 = 8)
// CHECK-SYSTEM: mul[0:5](2 * 4 = 8)
// CHECK-SYSTEM: mul[0:6](2 * 4 = 8)
// CHECK-SYSTEM: mul[0:7](2 * 4 = 8)
// arith.addf 8 + 4 = 12
// CHECK-SYSTEM: 8xf32=12 12 12 12 12 12 12 12

module @example {

  // Executable containing exported shims and calls to external functions.
  // Each executable can contain multiple exported functions and variants for
  // different architectures or even devices. It's also possible to mix hand-
  // authored functions with code generated ones even for the same functions
  // such that code generation is used as a fallback when the hand-authored
  // kernels aren't supported at runtime.
  stream.executable private @executable {
    stream.executable.export public @simple_mul workgroups(%workload: index) -> (index, index, index) {
      // This host function is used to compute the XYZ workgroup count
      // dispatched at runtime. It can query the %device for capabilities
      // and limits (last-level cache sizes, etc). The other arguments are the
      // values passed in the dispatch operation (usually things like root
      // output op tensor dimensions and other abstract values).
      %x = affine.apply affine_map<()[s0] -> (s0 ceildiv 64)>()[%workload]
      %c1 = arith.constant 1 : index
      stream.return %x, %c1, %c1 : index, index, index
    }

    builtin.module {
      // External function declaration using a user-chosen calling convention.
      func.func private @simple_mul_workgroup(
          %binding0: memref<f32>,
          %binding0_offset : index,
          %binding1: memref<f32>,
          %binding1_offset : index,
          %binding2: memref<f32>,
          %binding2_offset : index, %dim: index, %tid : index) attributes {
        // We can include some additional fields on the parameters struct as
        // needed. Here we request which processor is executing the call and
        // its data fields as defined by runtime/src/iree/schemas/cpu_data.h.
        hal.import.fields = ["processor_id", "processor_data"],
        llvm.bareptr = true
      }

      // IREE exported function using stream bindings and operands.
      // Compiler passes will be able to optimize across this interface and
      // deduplicate bindings/operands, convert/pack operands, and inline
      // constants operands.
      func.func @simple_mul(
          %binding0: !stream.binding,
          %binding1: !stream.binding,
          %binding2: !stream.binding,
          %dim: index) {
        %c0 = arith.constant 0 : index

        // This function is invoked once per workgroup so determine where this
        // particular workgroup is in the grid. In this example we use a
        // workgroup size of 64x1x1 (which is exceedingly small for CPUs but
        // useful for demonstration).
        %workgroup_id_x = flow.dispatch.workgroup.id[0] : index
        %tid = affine.apply affine_map<()[s0] -> (s0 * 64)>()[%workgroup_id_x]

        // Bindings are accessed by reference.
        %memref0 = stream.binding.subspan %binding0[%c0] : !stream.binding -> memref<?xf32>{%dim}
        %memref1 = stream.binding.subspan %binding1[%c0] : !stream.binding -> memref<?xf32>{%dim}
        %memref2 = stream.binding.subspan %binding2[%c0] : !stream.binding -> memref<?xf32>{%dim}

        %base0, %offset0, %size0, %stride0 = memref.extract_strided_metadata %memref0
            : memref<?xf32> -> memref<f32>, index, index, index
        %base1, %offset1, %size1, %stride1 = memref.extract_strided_metadata %memref1
            : memref<?xf32> -> memref<f32>, index, index, index
        %base2, %offset2, %size2, %stride2 = memref.extract_strided_metadata %memref2
            : memref<?xf32> -> memref<f32>, index, index, index
        

        // Call the externally defined C function with an (almost) plain C
        // calling convention. This will be fetched at runtime from the plugin binary.
        func.call @simple_mul_workgroup(
            %base0, %offset0, %base1, %offset1, %base2, %offset2, %dim, %workgroup_id_x)
            : (memref<f32>, index, memref<f32>, index, memref<f32>, index, index, index) -> ()

        // NOTE: this is code generated as normal - other MLIR ops can be used
        // here for looping/control flow, vector operations, linalg, etc.
        // This simple sample is just calling out to the external function but
        // microkernels fused with other code are possible.

        return
      }
    }
  }

  // Function demonstrating executable plugins and mixing plugins and codegen.
  // Invoke with:
  //  --device=local-sync
  //  --executable_plugin=system_plugin.so
  //  --function=mixed_invocation
  //  --input=8xf32=2
  //  --input=8xf32=4
  func.func @mixed_invocation(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) -> tensor<?xf32> {
    // The only externally available metadata in the dispatch are the values
    // passed in as operands. Here we pass in the dynamic dimension.
    %c0 = arith.constant 0 : index
    %dim = tensor.dim %arg0, %c0 : tensor<?xf32>

    // Dispatch a basic `ret = lhs * rhs` using an external function.
    // This form (@executable::@export) allows for automatic variant selection
    // when multi-targeting.
    %0 = flow.dispatch @executable::@simple_mul[%dim](%arg0, %arg1, %dim) : (tensor<?xf32>{%dim}, tensor<?xf32>{%dim}, index) -> tensor<?xf32>{%dim}

    // Code gen some other ops - these will interleave with hand-authored
    // ones but naturally won't be able to fuse with them.
    %1 = arith.addf %0, %arg1 : tensor<?xf32>

    return %1 : tensor<?xf32>
  }

}  // module